[New Model] support nemotron Audex (#4976)

Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

examples/offline_inference/audex/README.md

@@ -0,0 +1,114 @@
+# Audex (Nemotron-Labs-Audex-2B) offline inference
+
+One checkpoint, four tasks — each script builds its own `Omni` engine with
+the right deploy yaml, so they run as plain `python <script>` with no extra
+flags. Pass the HF repo ROOT as `--model` (default
+`nvidia/Nemotron-Labs-Audex-2B`); per-stage subfolders and download subsets
+resolve automatically.
+
+| script | pipeline (`vllm_omni/deploy/<name>.yaml`) | audio in | text out | speech out | general audio out |
+|---|---|---|---|---|---|
+| `text_to_speech.py` | `audex_tts` | ❌ | ❌ | ✅ | ❌ |
+| `text_to_audio.py` | `audex_tta` | ❌ | ❌ | ❌ | ✅ |
+| `audio_qa.py` | `audex_thinker_only` | ✅ | ✅ | ❌ | ❌ |
+| `speech_to_speech.py` | `audex_s2s` | ✅ | ✅ | ✅ | ❌ |
+
+Shared plumbing (corpus/WAV IO, tokenizer loading, the CFG cond/uncond
+pair contract) lives in `common.py`; each task script keeps only its own
+prompt recipe and flow. Note `text_to_audio.py` here drives the Audex
+two-stage AR pipeline — for diffusion text-to-audio models (Stable Audio
+Open, AudioX) use the generic
+`examples/offline_inference/text_to_audio/text_to_audio.py` instead.
+
+## text_to_speech.py — text → speech
+
+Thinker generates `<speechcodec_N>` tokens; the streaming causal decoder
+turns them into 16 kHz WAVs (one file per prompt, slugified filename).
+
+```bash
+# three built-in demo sentences -> results/audex_wavs/
+python examples/offline_inference/audex/text_to_speech.py
+
+# your own texts, with classifier-free guidance (official quality setting)
+python examples/offline_inference/audex/text_to_speech.py \
+    --texts "Hello there." "Nice to meet you." --cfg-scale 1.5 \
+    --output-dir results/my_tts
+```
+
+Key flags: `--texts` / `--texts-file` (TSV `utt_id<TAB>text`),
+`--cfg-scale` (default 1.0 = off; 1.5 recommended — guided requests run
+one at a time, each paired with a length-matched null prompt),
+`--temperature` (overrides the deploy yaml's stage-0 default).
+
+## text_to_audio.py — caption → general audio (sound effects)
+
+Same thinker over the interleaved 4-codebook `<audiocodec_N>` RVQ block,
+decoded by the external XCodec1 checkpoint. CFG is effectively mandatory
+(default scale 3.0). Clips are capped at `--codec-cap 4000` codec tokens
+(10 s).
+
+```bash
+# three built-in demo captions -> results/audex_tta_wavs/
+python examples/offline_inference/audex/text_to_audio.py
+
+python examples/offline_inference/audex/text_to_audio.py \
+    --captions "Thunder rolling across a valley." --output-dir results/my_tta
+```
+
+XCodec1 resolves from `--xcodec1-path`, the `XCODEC1_PATH` env var, or the
+default `hf-audio/xcodec-hubert-general-balanced` repo (downloaded on
+first use).
+
+## audio_qa.py — audio (+ instruction) → text
+
+Single-stage audio understanding on the full checkpoint (NV-Whisper
+encoder + projector + LM). The default question is ASR
+("Transcribe the input speech."); without `--audio-files`, vLLM's public
+`mary_had_lamb` asset is transcribed.
+
+```bash
+# transcribe the built-in demo asset
+python examples/offline_inference/audex/audio_qa.py
+
+# your own clips / free-form QA
+python examples/offline_inference/audex/audio_qa.py \
+    --audio-files a.wav b.wav --question "What language is being spoken?"
+```
+
+## speech_to_speech.py — spoken question → spoken answer
+
+The official three-pass cascade over ONE `audex_s2s` deployment: ASR
+(audio → transcript, text-final) → chat (transcript → answer, text-final)
+→ TTS (answer → speech through the streaming decoder). Only the TTS pass
+touches the codec path. Without `--audio-file`, the `mary_had_lamb` asset
+is used as the spoken input.
+
+```bash
+# full cascade on the built-in demo asset -> results/audex_s2s_answer.wav
+python examples/offline_inference/audex/speech_to_speech.py
+
+python examples/offline_inference/audex/speech_to_speech.py \
+    --audio-file question.wav --output results/answer.wav --cfg-scale 1.5
+```
+
+`--cfg-scale` (default 1.5) applies to the TTS pass only; 1.0 disables it.
+
+Online-serving counterparts (server + HTTP clients) live in
+`examples/online_serving/audex/`.
+
+## 30B-A3B (nvidia/Nemotron-Labs-Audex-30B-A3B)
+
+The same four scripts serve the 30B MoE checkpoint (hybrid Mamba + MoE
+NemotronH thinker; identical token space and decoder). The 30B REQUIRES an
+explicit 30B deploy yaml — the model root's default resolution lands on the
+2B-tuned configs:
+
+    python examples/offline_inference/audex/text_to_speech.py \
+        --model nvidia/Nemotron-Labs-Audex-30B-A3B \
+        --deploy-config vllm_omni/deploy/audex_tts_30b.yaml
+
+Per-mode yamls: `audex_{tts,tta,thinker_only,s2s}_30b.yaml`. Defaults fit a
+single H100 80 GB (thinker + decoder share the card; prefix caching is off —
+hybrid Mamba); if long sequences OOM, set `tensor_parallel_size: 2` on
+stage 0. First run downloads ~60 GB. Verified on a single H100 80 GB: ~61 GiB
+weights + healthy KV headroom; TP2 was not needed.

examples/offline_inference/audex/audio_qa.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) audio understanding example.
+
+Speech (or general audio) + text instruction → text, through the
+single-stage ``audex_thinker_only`` pipeline: NV-Whisper
+encoder + relu2 projector + 2B dense LM on ``checkpoint_folder_full``.
+
+The prompt format mirrors the official audioqa script: a ChatML user turn
+holding the ``<so_embedding>`` placeholder (expanded by the processor to
+750 embedding positions per 30 s clip) plus the instruction, with a closed
+``<think></think>`` priming.
+
+Examples (without --audio-files, vLLM's public ``mary_had_lamb`` asset is
+transcribed):
+
+    # Transcribe WAVs (ASR):
+    python examples/offline_inference/audex/audio_qa.py \\
+        --audio-files a.wav b.wav
+
+    # Free-form audio QA:
+    python examples/offline_inference/audex/audio_qa.py \\
+        --audio-files a.wav --question "What language is being spoken?"
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from common import default_deploy_config, load_audio_file, request_index
+from vllm.assets.audio import AudioAsset
+
+from vllm_omni import Omni
+
+ASR_QUESTION = "Transcribe the input speech."
+
+
+def build_prompt(question: str) -> str:
+    return f"<|im_start|>user\n<so_embedding>\n{question}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex audio understanding")
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument(
+        "--audio-files",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Input clips. Defaults to vLLM's mary_had_lamb audio asset.",
+    )
+    parser.add_argument("--question", type=str, default=ASR_QUESTION)
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        # The model root's default deploy yaml is the TTS pipeline; audio
+        # understanding needs the single-stage thinker-only pipeline.
+        default=default_deploy_config("audex_thinker_only.yaml"),
+        help="Deploy yaml (defaults to the audex_thinker_only pipeline).",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    prompt_text = build_prompt(args.question)
+    if args.audio_files:
+        clips = [load_audio_file(path) for path in args.audio_files]
+    else:
+        clips = [AudioAsset("mary_had_lamb").audio_and_sample_rate]
+    prompts = []
+    for audio, sr in clips:
+        prompts.append(
+            {
+                "prompt": prompt_text,
+                "multi_modal_data": {"audio": (audio, sr)},
+            }
+        )
+
+    outputs = engine.generate(prompts)
+
+    ordered = sorted(outputs, key=request_index)
+    names = [Path(p).name for p in args.audio_files] if args.audio_files else ["mary_had_lamb"]
+    for name, req_output in zip(names, ordered):
+        text = req_output.outputs[0].text if req_output.outputs else ""
+        print(f"{name}: {text.strip()}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/audex/common.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared helpers for the offline Audex examples in this directory.
+
+Import style follows the repo's example convention (same-directory module,
+see ``text_to_speech/qwen3_tts/tts_common.py``): scripts run as files, so
+``from common import ...`` resolves via the script's own directory.
+"""
+
+from __future__ import annotations
+
+import copy
+import re
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import soundfile as sf
+import torch
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_null_prompt
+
+SAMPLE_RATE = 16_000
+_DEPLOY_DIR = Path(__file__).resolve().parents[3] / "vllm_omni" / "deploy"
+
+# Guided decoding sharpens the distribution, so the unguided temperature
+# (0.1) adds excess sampling noise under CFG. Measured on the en-24 gate
+# corpus at cfg 1.5: temp 0.1 -> CER 7.31%, temp 0.05 -> CER 6.87% (vs the
+# unguided baseline 7.24%).
+GUIDED_TEMPERATURE = 0.05
+
+
+def default_deploy_config(name: str) -> str:
+    """Absolute path of a deploy yaml bundled with vllm_omni."""
+    return str(_DEPLOY_DIR / name)
+
+
+def slugify(text: str, fallback: str = "prompt") -> str:
+    slug = re.sub(r"\s+", "_", text.strip().lower())
+    slug = re.sub(r"[^a-z0-9_]+", "", slug)
+    return slug[:48] or fallback
+
+
+def load_corpus(
+    corpus_file: str | None,
+    texts: list[str] | None,
+    default_texts: tuple[str, ...],
+) -> list[tuple[str, str]]:
+    """(utt_id, text) pairs from a TSV corpus file, CLI texts, or defaults."""
+    if corpus_file:
+        corpus = []
+        for line in Path(corpus_file).read_text().splitlines():
+            if not line.strip():
+                continue
+            utt, text = line.split("\t", 1)
+            corpus.append((utt, text.strip()))
+        return corpus
+    resolved = texts if texts else list(default_texts)
+    return [(slugify(t), t) for t in resolved]
+
+
+def load_audio_file(path: str) -> tuple[np.ndarray, int]:
+    """Mono float32 samples + native sample rate (processors resample)."""
+    audio, sr = sf.read(path, dtype="float32")
+    if audio.ndim == 2:
+        audio = audio.mean(axis=1)
+    return audio, sr
+
+
+def extract_pcm(multimodal_output: dict) -> torch.Tensor:
+    """Flatten one request's audio payload to a 1-D float32 CPU tensor."""
+    audio = multimodal_output.get("model_outputs")
+    if audio is None:
+        audio = multimodal_output.get("audio")
+    if audio is None:
+        raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
+    if isinstance(audio, list):
+        valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
+        if not valid:
+            raise ValueError("audio list is empty")
+        return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
+    return torch.as_tensor(audio).float().cpu().reshape(-1)
+
+
+def write_wav(path: Path, pcm: torch.Tensor, sample_rate: int = SAMPLE_RATE) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    arr = (np.clip(pcm.numpy(), -1.0, 1.0) * 32767.0).astype(np.int16)
+    sf.write(str(path), arr, sample_rate, format="WAV", subtype="PCM_16")
+
+
+def load_audex_tokenizer(model: str, profile: str = "tts", folder: str = "checkpoint_folder_audiogen"):
+    """Thinker tokenizer for CFG/TTA prompt building (resolves the snapshot)."""
+    from transformers import AutoTokenizer
+
+    from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
+
+    root = ensure_audex_snapshot(model, profile=profile)
+    return AutoTokenizer.from_pretrained(str(Path(root) / folder), trust_remote_code=True)
+
+
+def request_index(req_output) -> int:
+    """Submission index from the numeric request-id prefix.
+
+    ``Omni.generate`` may return outputs out of submission order, and a
+    lexicographic sort would misplace "10_..." before "2_...".
+    """
+    match = re.search(r"(\d+)", str(req_output.request_id))
+    return int(match.group(1)) if match else 0
+
+
+def guided_stage_params(
+    engine: Omni,
+    cond_prompt: str,
+    tokenizer,
+    cfg_scale: float,
+    pair_id: str,
+    *,
+    temperature: float | None = None,
+    extra_args: dict[str, Any] | None = None,
+    null_prompt_builder: Callable[[str, Any], str] = build_null_prompt,
+) -> list:
+    """Stage sampling params carrying the CFG pair contract for ONE request.
+
+    Clones the engine's resolved defaults (never mutates the shared list),
+    applies ``extra_args`` unconditionally (e.g. the TTA RVQ contract), and
+    attaches the cond/uncond pair contract when ``cfg_scale > 1.0``. Guided
+    requests must go one at a time: each needs its own pair id and a
+    length-matched null prompt built from its own text.
+    """
+    params = copy.deepcopy(engine.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    if temperature is not None:
+        stage0.temperature = temperature
+    if stage0.extra_args is None:
+        stage0.extra_args = {}
+    if extra_args:
+        stage0.extra_args.update(extra_args)
+    if cfg_scale > 1.0:
+        stage0.extra_args.update(
+            {
+                "cfg_scale": float(cfg_scale),
+                "cfg_role": "cond",
+                "cfg_pair_id": pair_id,
+                "cfg_null_prompt": null_prompt_builder(cond_prompt, tokenizer),
+            }
+        )
+    return params

examples/offline_inference/audex/speech_to_speech.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) cascaded speech-to-speech example.
+
+The official three-pass cascade over ONE ``audex_s2s``
+deployment:
+
+  1. ASR pass  (audio + "Transcribe the input speech." → transcript, text)
+  2. chat pass (transcript → answer, text)
+  3. TTS pass  (answer → speech, audio via the streaming causal decoder)
+
+Text passes carry ``modalities: ["text"]`` and finish at stage 0; only the
+TTS pass (``modalities: ["audio"]``) streams codec frames into code2wav.
+
+Example (without --audio-file, vLLM's public ``mary_had_lamb`` asset is
+used as the spoken input):
+
+    python examples/offline_inference/audex/speech_to_speech.py \\
+        --audio-file question.wav --output results/answer.wav
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import re
+from pathlib import Path
+
+import torch
+from common import (
+    SAMPLE_RATE,
+    default_deploy_config,
+    guided_stage_params,
+    load_audex_tokenizer,
+    load_audio_file,
+    write_wav,
+)
+from vllm.assets.audio import AudioAsset
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+ASR_QUESTION = "Transcribe the input speech."
+SYSTEM_PROMPT = "You are a helpful and harmless assistant.\n\nYou are not allowed to use any tools."
+# The official chat-pass recipe below (SYSTEM_PROMPT + TEXT_INSTRUCTION
+# prefix, open <think> priming, temp 1.0/top_p 0.95, seed retry) applies to
+# the 30B-A3B ONLY (see is_30b()/_chat_prompt below and
+# tests/e2e/offline_inference/test_audex_s2s.py, which mirrors this same
+# split). It is load-bearing on the 30B: without the formatting-instruction
+# prefix a bare transcript turn is treated as a TTS request and the model
+# emits literal <speechcodec_N> tokens instead of a text answer (the system
+# prompt alone does not prevent this). The 2B must NOT use this recipe: it
+# is long-verified stable on a bare prompt + engine-default sampling, while
+# the official temp-1.0 seed-0 draw hallucinates off-language text (yet
+# closes </think> cleanly, so the seed retry never fires) and fails the
+# round-trip WER gate — measured, reproduced byte-identically.
+TEXT_INSTRUCTION = (
+    "SYSTEM & FORMATTING INSTRUCTION: You are Audex, created by NVIDIA based on the Nemotron-Cascade-2 architecture. "
+    "You may output your reasoning, followed by your final response. "
+    "You may format your reasoning block however you like. "
+    "However, your final response must strictly follow these rules: "
+    "* Write in plain, unformatted prose like a book or newspaper article. "
+    "* Do not use markdown, bullet points, lists, or headers. "
+    "* Standard numbers, abbreviations, and symbols are acceptable. "
+    "* [CRITICAL] You must press enter after every single sentence, placing each sentence on its own separate line."
+)
+
+
+def is_30b(model: str) -> bool:
+    # Users pass arbitrary repo IDs or local snapshot paths, so this is a
+    # heuristic (not an exact match against the official 30B-A3B repo ID)
+    # — good enough for an example script; keep in sync with the same
+    # branch in examples/online_serving/audex/client.py and
+    # tests/e2e/offline_inference/test_audex_s2s.py.
+    return "30b" in model.lower()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex cascaded S2S")
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument(
+        "--audio-file",
+        type=str,
+        default=None,
+        help="Spoken question (WAV). Defaults to vLLM's mary_had_lamb audio asset.",
+    )
+    parser.add_argument("--output", type=str, default="results/audex_s2s_answer.wav")
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        # The model root's default deploy yaml is the TTS-only pipeline; the
+        # cascade needs the audio-capable 2-stage full pipeline.
+        default=default_deploy_config("audex_s2s.yaml"),
+        help="Deploy yaml (defaults to the audex_s2s pipeline).",
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=1.5,
+        help="CFG strength for the TTS pass (1.0 disables; official setting 1.5).",
+    )
+    return parser.parse_args()
+
+
+def _chat_prompt(user_text: str) -> str:
+    # Official chat-pass recipe: system prompt + formatting instruction
+    # prefixed to the user text, thinking ENABLED via an open <think> block
+    # (the official demo's web UI ships "Enable reasoning" checked, i.e.
+    # build_chat_prompt(..., enable_thinking=True); chat_template.jinja then
+    # renders "<|im_start|>assistant\n<think>\n"). The instruction invites
+    # the model to reason before answering — priming the block open keeps
+    # that reasoning inside <think>...</think> so _extract_final_response
+    # below can drop it; with a closed block the 30B leaks the reasoning as
+    # plain text in front of the answer.
+    return (
+        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
+        f"<|im_start|>user\n{TEXT_INSTRUCTION}\n\n{user_text}<|im_end|>\n"
+        "<|im_start|>assistant\n<think>\n"
+    )
+
+
+def _asr_prompt() -> str:
+    return f"<|im_start|>user\n<so_embedding>\n{ASR_QUESTION}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+
+
+def _extract_final_response(text: str) -> str:
+    # Mirror the official demo's extract_spoken_answer
+    # (cascaded_s2s_web_server.py): keep only what follows the LAST closing
+    # think tag, falling back to the full text when the model skipped
+    # reasoning or never closed the block. Benign no-op for the 2B, whose
+    # answers carry no think block.
+    if "</think>" in text:
+        answer = text.rsplit("</think>", 1)[-1].strip()
+        if answer:
+            return answer
+    return text.strip()
+
+
+def _clean(text: str) -> str:
+    # The full checkpoint prefixes transcripts with a language tag sentence
+    # and quotes the content; extract the quoted payload when present. Span
+    # first-to-last quote (not split on every apostrophe) so contractions
+    # inside the transcript (e.g. "don't") survive intact.
+    text = text.strip()
+    match = re.search(r"'(.*)'", text, flags=re.S)
+    if match and match.group(1).strip():
+        return match.group(1).strip()
+    return text
+
+
+def _text_of(outputs) -> str:
+    (req_output,) = outputs
+    return (req_output.outputs[0].text or "").strip()
+
+
+CHAT_MAX_TOKENS = 2048  # official --text-max-tokens default
+# The official demo samples this pass unseeded; some trajectories deliberate
+# past the token budget without ever closing the think block, so retry a few
+# deterministic seeds until one closes (see the chat pass below).
+CHAT_SEEDS = (0, 1, 2)
+
+
+def _chat_params(engine: Omni, seed: int):
+    # Official demo text-pass sampling (cascaded_s2s_web_server.py defaults:
+    # --text-temperature 1.0, --text-top-p 0.95, --text-max-tokens 2048).
+    # The yaml stage-0 defaults (temperature 0.1, top_k 80) are TTS-thinker
+    # settings; under near-greedy decoding the 30B loops inside its
+    # reasoning block until the token budget runs out and never emits
+    # </think>, so the extracted "answer" would be truncated reasoning.
+    params = copy.deepcopy(engine.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    stage0.temperature = 1.0
+    stage0.top_p = 0.95
+    stage0.top_k = -1
+    stage0.max_tokens = CHAT_MAX_TOKENS
+    stage0.seed = seed
+    return params
+
+
+def _chat_once(engine: Omni, prompt: str, seed: int) -> tuple[str, bool]:
+    (req_output,) = engine.generate([{"prompt": prompt, "modalities": ["text"]}], _chat_params(engine, seed))
+    out = req_output.outputs[0]
+    text = (out.text or "").strip()
+    truncated = getattr(out, "finish_reason", None) == "length" or len(out.token_ids or []) >= CHAT_MAX_TOKENS
+    return text, truncated
+
+
+def main():
+    args = parse_args()
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    tokenizer = None
+    if args.cfg_scale > 1.0:
+        tokenizer = load_audex_tokenizer(args.model, profile="full", folder="checkpoint_folder_full")
+
+    if args.audio_file:
+        audio, sr = load_audio_file(args.audio_file)
+    else:
+        audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
+    # Pass 1 — ASR (text-final; never touches the speech decoder).
+    transcript = _clean(
+        _text_of(
+            engine.generate(
+                [
+                    {
+                        "prompt": _asr_prompt(),
+                        "multi_modal_data": {"audio": (audio, sr)},
+                        "modalities": ["text"],
+                    }
+                ]
+            )
+        )
+    )
+    print(f"[1/3] transcript : {transcript}")
+    if not transcript:
+        raise SystemExit("ASR pass produced an empty transcript; aborting cascade")
+
+    # Pass 2 — chat answer (text-final). The recipe is size-conditional (see
+    # the TEXT_INSTRUCTION comment above): the 30B requires the official
+    # recipe, the 2B regresses under it.
+    if is_30b(args.model):
+        # Official recipe. The raw completion is
+        # <reasoning...></think><answer>; speak only the final response.
+        # Retry with the next seed only when the budget ran out before the
+        # reasoning block closed (a natural stop without </think> is a
+        # complete answer); if every seed spins, fall back to the official
+        # demo's behavior of keeping the full text.
+        raw = ""
+        for seed in CHAT_SEEDS:
+            raw, truncated = _chat_once(engine, _chat_prompt(transcript), seed)
+            if "</think>" in raw or not truncated:
+                break
+            print(f"[2/3] chat seed {seed} hit the token budget mid-reasoning; retrying")
+        answer = _extract_final_response(raw)
+    else:
+        # 2B: bare prompt with a closed think block, engine-default sampling
+        # (yaml near-greedy), no retry; the answer is used as-is.
+        chat_prompt = f"<|im_start|>user\n{transcript}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+        answer = _text_of(engine.generate([{"prompt": chat_prompt, "modalities": ["text"]}]))
+    print(f"[2/3] answer     : {answer}")
+    if not answer:
+        raise SystemExit("Chat pass produced an empty answer; aborting cascade")
+
+    # Pass 3 — TTS (audio-final; streams through code2wav). Both final
+    # stages of the pipeline may emit an output record for the request;
+    # collect the audio-bearing one.
+    # The formatting instruction puts every sentence on its own line;
+    # flatten to plain prose for the TTS thinker.
+    tts_prompt = build_cond_prompt(re.sub(r"\s+", " ", answer).strip())
+    outputs = engine.generate(
+        [{"prompt": tts_prompt, "modalities": ["audio"]}],
+        guided_stage_params(engine, tts_prompt, tokenizer, args.cfg_scale, "s2s-tts-pass"),
+    )
+    chunks: list[torch.Tensor] = []
+    for req_output in outputs:
+        mm = getattr(req_output, "multimodal_output", None) or {}
+        if not mm and req_output.outputs:
+            mm = getattr(req_output.outputs[0], "multimodal_output", None) or {}
+        audio_val = mm.get("model_outputs", mm.get("audio"))
+        if audio_val is None:
+            continue
+        if isinstance(audio_val, list):
+            chunks.extend(torch.as_tensor(a).float().cpu().reshape(-1) for a in audio_val if a is not None)
+        else:
+            chunks.append(torch.as_tensor(audio_val).float().cpu().reshape(-1))
+    pcm = torch.cat(chunks) if chunks else torch.empty(0)
+    if pcm.numel() == 0:
+        raise SystemExit("TTS pass produced empty audio")
+
+    out_path = Path(args.output)
+    write_wav(out_path, pcm)
+    print(f"[3/3] speech     : {pcm.numel() / SAMPLE_RATE:.2f}s -> {out_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/audex/text_to_audio.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) text-to-audio inference example.
+
+Caption → general audio through the vLLM-Omni TTA pipeline: the audiogen
+thinker generates interleaved 4-codebook <audiocodec_N> RVQ tokens under an
+RVQ phase mask, and the external XCodec1 checkpoint decodes them to a 16 kHz
+mono WAV.
+
+(For DIFFUSION text-to-audio models — Stable Audio Open, AudioX — use
+``examples/offline_inference/text_to_audio/text_to_audio.py`` instead; Audex
+TTA is a two-stage autoregressive pipeline with a different interface.)
+
+Classifier-free guidance is effectively mandatory for TTA quality (official
+default scale 3.0), so every request submits a cond/uncond pair; requests
+run one at a time.
+
+Pass the HF repo ROOT as --model. XCodec1 resolves from --xcodec1-path /
+XCODEC1_PATH / the default hf-audio repo.
+
+Example:
+
+    python examples/offline_inference/audex/text_to_audio.py \\
+        --captions "Heavy rain falling on a tin roof." \\
+        --output-dir results/audex_tta_wavs
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from pathlib import Path
+
+from common import (
+    SAMPLE_RATE,
+    default_deploy_config,
+    extract_pcm,
+    guided_stage_params,
+    load_audex_tokenizer,
+    load_corpus,
+    write_wav,
+)
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_tta_cond_prompt, build_tta_null_prompt
+from vllm_omni.model_executor.models.audex.tta import build_tta_phase_token_ids
+
+DEFAULT_CAPTIONS = (
+    "Heavy rain falling on a tin roof.",
+    "A dog barking in the distance while birds chirp.",
+    "Ocean waves crashing on a rocky shore.",
+)
+DEFAULT_CODEC_CAP = 4000
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex TTA inference")
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument("--captions", type=str, nargs="+", default=None, help="Audio captions.")
+    parser.add_argument(
+        "--captions-file",
+        type=str,
+        default=None,
+        help="TSV corpus: one 'utt_id<TAB>caption' per line (overrides --captions).",
+    )
+    parser.add_argument("--output-dir", type=str, default="results/audex_tta_wavs")
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        # The model root's default deploy yaml is the TTS pipeline; TTA
+        # prompts and RVQ sampling params require the dedicated TTA pipeline.
+        default=default_deploy_config("audex_tta.yaml"),
+        help="Deploy yaml (defaults to the audex_tta pipeline).",
+    )
+    parser.add_argument(
+        "--xcodec1-path",
+        type=str,
+        default=None,
+        help="Local XCodec1 checkpoint (defaults to $XCODEC1_PATH or the hf-audio repo).",
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=3.0,
+        help="Classifier-free guidance strength (official TTA default 3.0; 1.0 disables).",
+    )
+    parser.add_argument(
+        "--codec-cap",
+        type=int,
+        default=DEFAULT_CODEC_CAP,
+        help="Max generated codec tokens before the phase mask forces <audiogen_end>.",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    if args.xcodec1_path:
+        os.environ["XCODEC1_PATH"] = args.xcodec1_path
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    corpus = load_corpus(args.captions_file, args.captions, DEFAULT_CAPTIONS)
+
+    tokenizer = load_audex_tokenizer(args.model, profile="tta")
+    phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer)
+    tta_rvq = {
+        "phase_token_ids": phase_token_ids,
+        "start_tid": start_tid,
+        "end_tid": end_tid,
+        "codec_cap": args.codec_cap,
+        # The prompt ends with <audiogen_start>, so generation begins at
+        # RVQ phase 0 immediately.
+        "start_in_prompt": True,
+    }
+
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    print(f"Model       : {args.model}")
+    print(f"Captions    : {len(corpus)}")
+    print(f"CFG scale   : {args.cfg_scale}")
+    print(f"Codec cap   : {args.codec_cap}")
+    print(f"Output dir  : {output_dir}")
+
+    total_elapsed = 0.0
+    total_dur = 0.0
+    for utt, caption in corpus:
+        cond_prompt = build_tta_cond_prompt(caption)
+        sampling = guided_stage_params(
+            engine,
+            cond_prompt,
+            tokenizer,
+            args.cfg_scale,
+            f"tta-{utt}",
+            extra_args={"tta_rvq": tta_rvq},
+            null_prompt_builder=build_tta_null_prompt,
+        )
+        t_start = time.perf_counter()
+        outputs = engine.generate([cond_prompt], sampling)
+        total_elapsed += time.perf_counter() - t_start
+
+        (req_output,) = outputs
+        pcm = extract_pcm(req_output.outputs[0].multimodal_output)
+        dur = pcm.numel() / SAMPLE_RATE
+        if dur <= 0:
+            raise RuntimeError(f"empty audio for {utt}")
+        out_path = output_dir / f"{utt}.wav"
+        write_wav(out_path, pcm)
+        total_dur += dur
+        print(f"  {utt:<48} dur={dur:6.2f}s  -> {out_path}")
+
+    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
+    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/audex/text_to_speech.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) TTS inference example.
+
+Runs Stage 0 (thinker) + Stage 1 (streaming causal speech decoder) end-to-end
+through the vLLM-Omni engine and saves a 16 kHz mono WAV per prompt.
+
+Pass the HF repo ROOT as --model; per-stage subfolders resolve automatically.
+
+Examples:
+
+    python examples/offline_inference/audex/text_to_speech.py \\
+        --texts "The weather is so good today." \\
+        --output-dir results/audex_wavs
+
+    # TSV corpus (utt_id<TAB>text per line), all requests in flight at once:
+    python examples/offline_inference/audex/text_to_speech.py \\
+        --texts-file corpus.tsv --output-dir results/audex_wavs
+
+    # One request at a time:
+    python examples/offline_inference/audex/text_to_speech.py \\
+        --texts-file corpus.tsv --sequential --output-dir results/audex_wavs
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+from common import (
+    GUIDED_TEMPERATURE,
+    SAMPLE_RATE,
+    extract_pcm,
+    guided_stage_params,
+    load_audex_tokenizer,
+    load_corpus,
+    request_index,
+    write_wav,
+)
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+DEFAULT_TEXTS = (
+    "Hello world.",
+    "The quick brown fox jumps over the lazy dog.",
+    "Today is a beautiful day for a walk in the park.",
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex TTS inference")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="nvidia/Nemotron-Labs-Audex-2B",
+        help="HF repo root (or local snapshot root).",
+    )
+    parser.add_argument("--texts", type=str, nargs="+", default=None, help="Plain-text prompts.")
+    parser.add_argument(
+        "--texts-file",
+        type=str,
+        default=None,
+        help="TSV corpus: one 'utt_id<TAB>text' per line (overrides --texts).",
+    )
+    parser.add_argument("--output-dir", type=str, default="results/audex_wavs")
+    parser.add_argument("--deploy-config", type=str, default=None, help="Override the deploy config path.")
+    parser.add_argument(
+        "--sequential",
+        action="store_true",
+        help="Submit one request at a time instead of all at once.",
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=1.0,
+        help=(
+            "Classifier-free guidance strength. 1.0 disables guidance (default, "
+            "matches the v1 baseline); 1.5 is the official quality setting. "
+            "Guided runs submit one request at a time (each carries its own "
+            "length-matched null prompt and pair id)."
+        ),
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=None,
+        help="Override the deploy yaml's stage-0 sampling temperature.",
+    )
+    return parser.parse_args()
+
+
+def _write_outputs(outputs, batch: list[tuple[str, str]], output_dir: Path) -> float:
+    ordered = sorted(outputs, key=request_index)
+    if len(ordered) != len(batch):
+        raise RuntimeError(f"expected {len(batch)} outputs, got {len(ordered)}")
+
+    total_dur = 0.0
+    for (utt, _text), req_output in zip(batch, ordered):
+        pcm = extract_pcm(req_output.outputs[0].multimodal_output)
+        dur = pcm.numel() / SAMPLE_RATE
+        if dur <= 0:
+            raise RuntimeError(f"empty audio for {utt}")
+        out_path = output_dir / f"{utt}.wav"
+        write_wav(out_path, pcm)
+        total_dur += dur
+        print(f"  {utt:<40} dur={dur:6.2f}s  -> {out_path}")
+    return total_dur
+
+
+def main():
+    args = parse_args()
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    corpus = load_corpus(args.texts_file, args.texts, DEFAULT_TEXTS)
+
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    cfg_enabled = args.cfg_scale > 1.0
+    tokenizer = load_audex_tokenizer(args.model) if cfg_enabled else None
+
+    print(f"Model       : {args.model}")
+    print(f"Prompts     : {len(corpus)}")
+    print(f"Mode        : {'sequential' if args.sequential or cfg_enabled else 'batched'}")
+    print(f"CFG scale   : {args.cfg_scale}" + (" (guidance off)" if not cfg_enabled else ""))
+    print(f"Output dir  : {output_dir}")
+
+    total_elapsed = 0.0
+    total_dur = 0.0
+    # CFG runs one request at a time: each request needs its own pair id and
+    # text-specific null prompt in the stage-0 sampling params.
+    batches = [[item] for item in corpus] if (args.sequential or cfg_enabled) else [corpus]
+    for batch in batches:
+        prompts = [build_cond_prompt(text) for _utt, text in batch]
+        sampling_params_list = None
+        if cfg_enabled:
+            sampling_params_list = guided_stage_params(
+                engine,
+                prompts[0],
+                tokenizer,
+                args.cfg_scale,
+                f"cfg-{batch[0][0]}",
+                temperature=GUIDED_TEMPERATURE,
+            )
+        if args.temperature is not None:
+            import copy
+
+            if sampling_params_list is None:
+                sampling_params_list = copy.deepcopy(engine.resolve_sampling_params_list(None))
+            sampling_params_list[0].temperature = args.temperature
+        t_start = time.perf_counter()
+        outputs = engine.generate(prompts, sampling_params_list)
+        total_elapsed += time.perf_counter() - t_start
+        total_dur += _write_outputs(outputs, batch, output_dir)
+
+    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
+    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/audex/README.md

@@ -0,0 +1,89 @@
+# Audex (Nemotron-Labs-Audex-2B) online serving
+
+One checkpoint, four deployment modes. `run_server.sh` starts the server
+for a given `MODE`; `client.py --mode <mode>` tests it. The capability
+matrix:
+
+| mode (`vllm_omni/deploy/audex_<mode>.yaml`) | audio in | text out | speech out | general audio out | endpoint |
+|---|---|---|---|---|---|
+| `tts` | ❌ | ❌ | ✅ | ❌ | `/v1/audio/speech` |
+| `tta` | ❌ | ❌ | ❌ | ✅ | `/v1/audio/speech` |
+| `thinker_only` | ✅ | ✅ | ❌ | ❌ | `/v1/chat/completions` |
+| `s2s` | ✅ | ✅ | ✅ | ❌ | both |
+
+Audio-input modes of `client.py` fall back to vLLM's public
+`mary_had_lamb` asset when `--audio-file` is omitted; CFG defaults follow
+the official settings per mode (tts/s2s 1.5, tta 3.0; `--cfg-scale 1.0`
+disables).
+
+## tts — text → speech
+
+Text-only thinker (`checkpoint_folder_audiogen`) emits `<speechcodec_N>`
+tokens; the streaming causal decoder produces 16 kHz WAVs. This is also
+the DEFAULT pipeline when the repo root is served without a deploy config.
+
+```bash
+./run_server.sh                       # MODE=tts, port 8097
+
+python client.py --mode tts --text "Hello there." --output hello.wav
+# or raw curl:
+curl -s http://localhost:8097/v1/audio/speech \
+    -H 'Content-Type: application/json' \
+    -d '{"model": "nvidia/Nemotron-Labs-Audex-2B", "input": "Hello there.", "response_format": "wav", "extra_params": {"cfg_scale": 1.5}}' \
+    -o hello.wav
+```
+
+## tta — caption → general audio
+
+Same thinker over the 4-codebook `<audiocodec_N>` RVQ block, decoded by
+the external XCodec1 checkpoint (auto-downloaded; override with
+`XCODEC1_PATH`). Clips are capped at 10 s.
+
+```bash
+MODE=tta ./run_server.sh
+
+python client.py --mode tta --caption "Heavy rain falling on a tin roof." \
+    --output rain.wav
+```
+
+## thinker_only — audio (+ instruction) → text
+
+Single-stage audio understanding on the full checkpoint (NV-Whisper
+encoder + projector + LM). Send the clip as `input_audio` chat content.
+
+```bash
+MODE=thinker_only ./run_server.sh
+
+python client.py --mode thinker_only                      # transcribe the demo asset
+python client.py --mode thinker_only --audio-file a.wav \
+    --question "What language is being spoken?"
+```
+
+## s2s — spoken question → spoken answer
+
+The audio-capable thinker plus the speech decoder in one deployment.
+`client.py --mode s2s` runs the official three passes: ASR (chat with
+`input_audio`) → chat on the transcript → TTS on the answer. Text-final
+passes carry `"modalities": ["text"]`; only the TTS pass streams through
+the decoder.
+
+```bash
+MODE=s2s PORT=8098 ./run_server.sh
+
+python client.py --mode s2s --port 8098 \
+    --audio-file question.wav --output answer.wav
+```
+
+Offline counterparts (no HTTP, one `Omni` engine per script) live in
+`examples/offline_inference/audex/`.
+
+## 30B-A3B (nvidia/Nemotron-Labs-Audex-30B-A3B)
+
+`SIZE=30b` switches run_server.sh to the 30B yamls and model id:
+
+    SIZE=30b ./run_server.sh                 # 30B tts on port 8097
+    SIZE=30b MODE=s2s PORT=8098 ./run_server.sh
+
+Pass `--model nvidia/Nemotron-Labs-Audex-30B-A3B` to the client — the server
+validates the model id and 404s on a mismatch. Single-H100 defaults; see the
+yaml comments for the TP2 fallback. First launch downloads ~60 GB.

examples/online_serving/audex/client.py

@@ -0,0 +1,319 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Audex online client for all four deployment modes.
+
+Pick ``--mode`` to match the server's deploy yaml (see run_server.sh):
+
+  tts           /v1/audio/speech: text -> speech WAV
+  tta           /v1/audio/speech: caption -> general-audio WAV
+  thinker_only  /v1/chat/completions: audio (+ question) -> text
+  s2s           the official three-pass cascade against one audex_s2s
+                server: ASR (chat, input_audio) -> chat -> TTS
+
+Audio-input modes fall back to vLLM's public ``mary_had_lamb`` asset when
+``--audio-file`` is omitted.
+
+Examples:
+
+    python examples/online_serving/audex/client.py --mode tts \\
+        --text "Hello there." --output hello.wav
+    python examples/online_serving/audex/client.py --mode tta \\
+        --caption "Heavy rain falling on a tin roof." --output rain.wav
+    python examples/online_serving/audex/client.py --mode thinker_only
+    python examples/online_serving/audex/client.py --mode s2s \\
+        --audio-file question.wav --output answer.wav
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import re
+from pathlib import Path
+
+import requests
+
+ASR_QUESTION = "Transcribe the input speech."
+SYSTEM_PROMPT = "You are a helpful and harmless assistant.\n\nYou are not allowed to use any tools."
+# The official chat-pass recipe below (SYSTEM_PROMPT + TEXT_INSTRUCTION
+# prefix, reasoning enabled, temp 1.0/top_p 0.95, seed retry) applies to the
+# 30B-A3B ONLY (see is_30b() below and
+# tests/e2e/offline_inference/test_audex_s2s.py, which mirrors this same
+# split). It is load-bearing on the 30B: without the formatting-instruction
+# prefix a bare transcript turn is treated as a TTS request and the model
+# emits literal <speechcodec_N> tokens instead of a text answer (the system
+# prompt alone does not prevent this). The 2B must NOT use this recipe: the
+# official temp-1.0 seed-0 draw measurably hallucinates off-language text on
+# the 2B (yet closes </think> cleanly, so the seed retry never fires) and
+# fails the round-trip WER gate — measured, reproduced byte-identically. The
+# 2B instead keeps its original bare-message / enable_thinking=False /
+# temperature=0.0 recipe.
+TEXT_INSTRUCTION = (
+    "SYSTEM & FORMATTING INSTRUCTION: You are Audex, created by NVIDIA based on the Nemotron-Cascade-2 architecture. "
+    "You may output your reasoning, followed by your final response. "
+    "You may format your reasoning block however you like. "
+    "However, your final response must strictly follow these rules: "
+    "* Write in plain, unformatted prose like a book or newspaper article. "
+    "* Do not use markdown, bullet points, lists, or headers. "
+    "* Standard numbers, abbreviations, and symbols are acceptable. "
+    "* [CRITICAL] You must press enter after every single sentence, placing each sentence on its own separate line."
+)
+DEFAULT_TEXT = "Hello there! This is the Audex text to speech pipeline."
+DEFAULT_CAPTION = "Heavy rain falling on a tin roof."
+# Official quality settings; 1.0 disables guidance.
+DEFAULT_CFG = {"tts": 1.5, "tta": 3.0, "s2s": 1.5}
+
+
+def is_30b(model: str) -> bool:
+    # Users pass arbitrary repo IDs or local snapshot paths, so this is a
+    # heuristic (not an exact match against the official 30B-A3B repo ID)
+    # — good enough for an example script; keep in sync with the same
+    # branch in examples/offline_inference/audex/speech_to_speech.py and
+    # tests/e2e/offline_inference/test_audex_s2s.py.
+    return "30b" in model.lower()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Audex online client (tts / tta / thinker_only / s2s)")
+    parser.add_argument("--mode", choices=("tts", "tta", "thinker_only", "s2s"), default="s2s")
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8097)
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument("--text", type=str, default=DEFAULT_TEXT, help="TTS input text (mode=tts).")
+    parser.add_argument("--caption", type=str, default=DEFAULT_CAPTION, help="TTA caption (mode=tta).")
+    parser.add_argument(
+        "--audio-file",
+        type=str,
+        default=None,
+        help="Input WAV for thinker_only/s2s. Defaults to vLLM's mary_had_lamb asset.",
+    )
+    parser.add_argument("--question", type=str, default=ASR_QUESTION, help="Instruction (mode=thinker_only).")
+    parser.add_argument("--output", type=str, default="results/audex_online.wav")
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=None,
+        help="CFG strength; defaults per mode (tts/s2s 1.5, tta 3.0). 1.0 disables.",
+    )
+    parser.add_argument("--max-tokens", type=int, default=512)
+    return parser.parse_args()
+
+
+def _audio_b64(audio_file: str | None) -> str:
+    if audio_file:
+        return base64.b64encode(Path(audio_file).read_bytes()).decode()
+    # Fall back to vLLM's public asset, re-encoded as WAV for the chat API.
+    import io
+
+    import numpy as np
+    import soundfile as sf
+    from vllm.assets.audio import AudioAsset
+
+    audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
+    buf = io.BytesIO()
+    sf.write(buf, np.asarray(audio, dtype="float32"), int(sr), format="WAV")
+    return base64.b64encode(buf.getvalue()).decode()
+
+
+def _chat(
+    base_url: str,
+    model: str,
+    messages: list[dict],
+    max_tokens: int,
+    enable_thinking: bool = False,
+    temperature: float = 0.0,
+    top_p: float = 1.0,
+    seed: int = 0,
+) -> tuple[str, str]:
+    resp = requests.post(
+        f"{base_url}/v1/chat/completions",
+        json={
+            "model": model,
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "seed": seed,
+            # Text-final passes must say so explicitly: without "modalities"
+            # the server falls back to the deployment's configured output
+            # modalities, which for the s2s pipeline include the audio final
+            # stage — routing ASR/chat through code2wav instead of stopping
+            # at stage 0.
+            "modalities": ["text"],
+            # The checkpoint's chat template thinking toggle. The official
+            # cascaded demo transcribes with enable_thinking=False but runs
+            # the chat pass with reasoning ENABLED (its web UI ships the
+            # "Enable reasoning" checkbox checked), priming the generation
+            # prompt with an open "<think>\n" so the model's reasoning stays
+            # inside the think block; _extract_final_response then drops it.
+            # With thinking disabled the 30B leaks its reasoning as plain
+            # text in front of the final response.
+            "chat_template_kwargs": {"enable_thinking": enable_thinking},
+        },
+        timeout=300,
+    )
+    resp.raise_for_status()
+    choice = resp.json()["choices"][0]
+    return choice["message"]["content"].strip(), choice.get("finish_reason") or ""
+
+
+def _speech(base_url: str, model: str, text: str, cfg_scale: float, out_path: Path) -> None:
+    payload: dict = {"model": model, "input": text, "response_format": "wav"}
+    if cfg_scale > 1.0:
+        payload["extra_params"] = {"cfg_scale": cfg_scale}
+    resp = requests.post(f"{base_url}/v1/audio/speech", json=payload, timeout=600)
+    resp.raise_for_status()
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_bytes(resp.content)
+
+
+def _audio_question(audio_b64: str, question: str) -> list[dict]:
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "input_audio", "input_audio": {"data": audio_b64, "format": "wav"}},
+                {"type": "text", "text": question},
+            ],
+        }
+    ]
+
+
+def _tts_text(answer: str) -> str:
+    # The formatting instruction's CRITICAL rule puts each final-response
+    # sentence on its own line, and the official demo TTS-es those sentence
+    # segments (split_sentence_segments). Mirror that per-line segmentation
+    # but keep a single /v1/audio/speech request: speak the leading whole
+    # sentences of the CLEANED final response up to ~200 chars, never
+    # cutting mid-sentence. Markdown markup is stripped as before (it makes
+    # the TTS thinker stop almost immediately).
+    sentences = []
+    for line in answer.splitlines():
+        line = re.sub(r"[*#_`]+", "", line)
+        line = re.sub(r"\s+", " ", line).strip()
+        if line:
+            sentences.append(line)
+    picked: list[str] = []
+    total = 0
+    for sentence in sentences:
+        if picked and total + len(sentence) > 200:
+            break
+        picked.append(sentence)
+        total += len(sentence) + 1
+    return " ".join(picked)
+
+
+def _extract_final_response(text: str) -> str:
+    # Mirror the official demo's extract_spoken_answer
+    # (cascaded_s2s_web_server.py): the chat pass reasons inside the primed
+    # <think> block, so keep only what follows the LAST closing tag; fall
+    # back to the full text when the model skipped reasoning or never
+    # closed the block. Benign no-op for the 2B's clean answers.
+    if "</think>" in text:
+        answer = text.rsplit("</think>", 1)[-1].strip()
+        if answer:
+            return answer
+    return text.strip()
+
+
+def _clean_transcript(text: str) -> str:
+    # The full checkpoint prefixes transcripts with a language-tag sentence
+    # and quotes the content; extract the quoted payload when present. Span
+    # first-to-last quote (not split on every apostrophe) so contractions
+    # inside the transcript (e.g. "don't") survive intact.
+    text = text.strip()
+    match = re.search(r"'(.*)'", text, flags=re.S)
+    if match and match.group(1).strip():
+        return match.group(1).strip()
+    return text
+
+
+def main():
+    args = parse_args()
+    base_url = f"http://{args.host}:{args.port}"
+    cfg_scale = args.cfg_scale if args.cfg_scale is not None else DEFAULT_CFG.get(args.mode, 1.0)
+    out_path = Path(args.output)
+
+    if args.mode == "tts":
+        _speech(base_url, args.model, args.text, cfg_scale, out_path)
+        print(f"speech: {out_path.stat().st_size} bytes -> {out_path}")
+        return
+
+    if args.mode == "tta":
+        _speech(base_url, args.model, args.caption, cfg_scale, out_path)
+        print(f"audio : {out_path.stat().st_size} bytes -> {out_path}")
+        return
+
+    audio_b64 = _audio_b64(args.audio_file)
+
+    if args.mode == "thinker_only":
+        content, _ = _chat(base_url, args.model, _audio_question(audio_b64, args.question), args.max_tokens)
+        print(f"answer: {_extract_final_response(content)}")
+        return
+
+    # s2s — the official three-pass cascade.
+    asr_content, _ = _chat(base_url, args.model, _audio_question(audio_b64, ASR_QUESTION), args.max_tokens)
+    transcript = _clean_transcript(asr_content)
+    print(f"[1/3] transcript : {transcript}")
+    if not transcript:
+        raise SystemExit("ASR pass produced an empty transcript; aborting cascade")
+
+    # The chat-pass recipe is size-conditional (see the TEXT_INSTRUCTION
+    # comment above): the 30B requires the official recipe, the 2B keeps its
+    # original bare-message recipe.
+    if is_30b(args.model):
+        # Official recipe: system prompt + formatting instruction prefixed
+        # to the transcript, reasoning enabled, and the official demo's
+        # text-pass sampling (its defaults: --text-temperature 1.0,
+        # --text-top-p 0.95, --text-max-tokens 2048). Near-greedy decoding
+        # makes the 30B loop inside its reasoning block until the token
+        # budget runs out without ever closing </think>. Even at the
+        # official temperature some trajectories spin past the budget, so
+        # retry a few deterministic seeds; a natural stop without </think>
+        # is a complete answer, and if every seed spins we keep the full
+        # text like the official demo does.
+        chat_messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": f"{TEXT_INSTRUCTION}\n\n{transcript}"},
+        ]
+        raw = ""
+        for seed in (0, 1, 2):
+            raw, finish_reason = _chat(
+                base_url,
+                args.model,
+                chat_messages,
+                max(args.max_tokens, 2048),
+                enable_thinking=True,
+                temperature=1.0,
+                top_p=0.95,
+                seed=seed,
+            )
+            if "</think>" in raw or finish_reason != "length":
+                break
+            print(f"[2/3] chat seed {seed} hit the token budget mid-reasoning; retrying")
+        answer = _extract_final_response(raw)
+    else:
+        # 2B: bare user message (no system/instruction prefix), reasoning
+        # disabled, temperature 0.0, no retry — its original recipe. The
+        # official temp-1.0 seed-0 recipe measurably makes the 2B
+        # hallucinate off-language answers (see the module-level comment).
+        chat_messages = [{"role": "user", "content": transcript}]
+        raw, _ = _chat(
+            base_url,
+            args.model,
+            chat_messages,
+            args.max_tokens,
+            enable_thinking=False,
+            temperature=0.0,
+        )
+        answer = _extract_final_response(raw)
+    print(f"[2/3] answer     : {answer}")
+    if not answer:
+        raise SystemExit("Chat pass produced an empty answer; aborting cascade")
+
+    _speech(base_url, args.model, _tts_text(answer), cfg_scale, out_path)
+    print(f"[3/3] speech     : {out_path.stat().st_size} bytes -> {out_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/audex/run_server.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+# Launch the vLLM-Omni server for Audex (Nemotron-Labs-Audex-2B).
+#
+# MODE picks the deployment (see README.md for the capability matrix):
+#   tts           text -> speech            /v1/audio/speech   (default)
+#   tta           caption -> general audio  /v1/audio/speech
+#   thinker_only  audio -> text             /v1/chat/completions
+#   s2s           cascaded speech-to-speech (both endpoints)
+#
+# Pass the HF repo ROOT as MODEL — per-stage subfolders resolve
+# automatically.
+#
+# Usage:
+#   ./run_server.sh                       # 2B tts on port 8097, GPU 0
+#   SIZE=30b ./run_server.sh              # 30B-A3B tts (explicit 30B yaml)
+#   SIZE=30b MODE=s2s PORT=8098 ./run_server.sh
+#   MODE=tta ./run_server.sh              # needs XCodec1 (auto-downloaded)
+#   MODEL=/path/to/local/snapshot ./run_server.sh
+
+set -e
+
+MODE="${MODE:-${1:-tts}}"
+SIZE="${SIZE:-2b}"
+if [ "$SIZE" = "30b" ]; then
+    MODEL="${MODEL:-nvidia/Nemotron-Labs-Audex-30B-A3B}"
+    YAML_SUFFIX="_30b"
+elif [ "$SIZE" = "2b" ]; then
+    MODEL="${MODEL:-nvidia/Nemotron-Labs-Audex-2B}"
+    YAML_SUFFIX=""
+else
+    echo "Unknown SIZE '$SIZE'; expected 2b|30b" >&2
+    exit 1
+fi
+PORT="${PORT:-8097}"
+GPUS="${GPUS:-0}"
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
+DEPLOY_YAML="$REPO_ROOT/vllm_omni/deploy/audex_${MODE}${YAML_SUFFIX}.yaml"
+if [ ! -f "$DEPLOY_YAML" ]; then
+    echo "Unknown MODE '$MODE' / SIZE '$SIZE' (no $DEPLOY_YAML); expected MODE=tts|tta|thinker_only|s2s, SIZE=2b|30b" >&2
+    exit 1
+fi
+
+echo "Starting Audex server"
+echo "  MODE=$MODE ($DEPLOY_YAML)"
+echo "  SIZE=$SIZE"
+echo "  MODEL=$MODEL"
+echo "  PORT=$PORT"
+echo "  CUDA_VISIBLE_DEVICES=$GPUS"
+
+CUDA_VISIBLE_DEVICES="$GPUS" \
+vllm-omni serve "$MODEL" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --trust-remote-code \
+    --stage-configs-path "$DEPLOY_YAML" \
+    --omni

recipes/NVIDIA/Nemotron-Labs-Audex.md

@@ -0,0 +1,142 @@
+# Nemotron-Labs Audex (2B / 30B-A3B)
+
+> Speech suite on 1×H100: TTS, text-to-audio, audio understanding, and cascaded speech-to-speech
+
+## Summary
+
+- Vendor: NVIDIA
+- Model: `nvidia/Nemotron-Labs-Audex-2B`, `nvidia/Nemotron-Labs-Audex-30B-A3B`
+- Task: Text-to-speech (English), text-to-audio (sound effects), audio understanding (ASR / audio QA), cascaded speech-to-speech
+- Mode: Online serving (`/v1/audio/speech`, `/v1/chat/completions`) and offline batch inference
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to serve either Audex checkpoint through one of its four
+deployment pipelines. Both sizes share the same repo layout, codec token
+space, and streaming causal speech decoder (16 kHz output); they differ in
+the thinker: the 2B is a dense LM, the 30B-A3B is a hybrid Mamba + MoE
+NemotronH (~3B active parameters). Pick the pipeline by task:
+
+| pipeline (`vllm_omni/deploy/<name>[_30b].yaml`) | audio in | text out | speech out | general audio out | endpoint |
+|---|---|---|---|---|---|
+| `audex_tts` | ❌ | ❌ | ✅ | ❌ | `/v1/audio/speech` |
+| `audex_tta` | ❌ | ❌ | ❌ | ✅ | `/v1/audio/speech` |
+| `audex_thinker_only` | ✅ | ✅ | ❌ | ❌ | `/v1/chat/completions` |
+| `audex_s2s` | ✅ | ✅ | ✅ | ❌ | both |
+
+## References
+
+- Model cards: [Nemotron-Labs-Audex-2B](https://huggingface.co/nvidia/Nemotron-Labs-Audex-2B), [Nemotron-Labs-Audex-30B-A3B](https://huggingface.co/nvidia/Nemotron-Labs-Audex-30B-A3B)
+- Offline examples: [`examples/offline_inference/audex/`](../../examples/offline_inference/audex/)
+- Online examples: [`examples/online_serving/audex/`](../../examples/online_serving/audex/)
+
+## Hardware Support
+
+## GPU
+
+### 1×H100 80GB — Audex-2B
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12+
+- CUDA: 12.x
+- vLLM version: 0.24.0
+- vLLM-Omni version or commit: PR #4976 (branch `audex`)
+
+#### Command
+
+**Online serving** (TTS is the default pipeline; the repo root's
+`model_type` auto-resolves `audex_tts.yaml`):
+
+```bash
+vllm-omni serve nvidia/Nemotron-Labs-Audex-2B \
+    --host 0.0.0.0 --port 8097 \
+    --trust-remote-code --omni
+# other modes: add --stage-configs-path vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml
+# or use the launcher: MODE=s2s PORT=8098 examples/online_serving/audex/run_server.sh
+```
+
+**Offline batch inference** (one script per task, each defaults to its
+correct deploy yaml):
+
+```bash
+python examples/offline_inference/audex/text_to_speech.py \
+    --texts "Hello world." --cfg-scale 1.5 --output-dir results/audex_wavs
+python examples/offline_inference/audex/text_to_audio.py     # sound effects
+python examples/offline_inference/audex/audio_qa.py          # ASR / audio QA
+python examples/offline_inference/audex/speech_to_speech.py  # 3-pass cascade
+```
+
+#### Verification
+
+```bash
+curl -X POST http://localhost:8097/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "nvidia/Nemotron-Labs-Audex-2B",
+        "input": "Hello, how are you?",
+        "response_format": "wav",
+        "extra_params": {"cfg_scale": 1.5}
+    }' --output hello.wav
+# or the universal client (all four modes):
+python examples/online_serving/audex/client.py --mode tts --output hello.wav
+```
+
+#### Notes
+
+- Memory usage: TTS deploy — stage 0 (thinker) 0.4, stage 1 (decoder) 0.25
+  GPU-memory fraction; both stages share GPU 0.
+- Key flags: `--trust-remote-code` and `--omni` are required.
+- Classifier-free guidance: `extra_params.cfg_scale` (1.0 = off; 1.5 is the
+  official quality setting; TTA effectively requires its default 3.0).
+  Measured en-24 self-transcribed CER: 6.87% guided vs 7.24% unguided.
+- Output: 16 kHz mono WAV. English, single built-in voice (no cloning).
+- TTA decodes through the external XCodec1 checkpoint
+  (`hf-audio/xcodec-hubert-general-balanced`, auto-downloaded; override
+  with `XCODEC1_PATH`).
+
+### 1×H100 80GB — Audex-30B-A3B
+
+#### Environment
+
+Same as the 2B row above.
+
+#### Command
+
+The 30B REQUIRES its explicit deploy yaml (the shared HF `model_type`
+auto-resolves to the 2B-tuned config otherwise):
+
+```bash
+vllm-omni serve nvidia/Nemotron-Labs-Audex-30B-A3B \
+    --host 0.0.0.0 --port 8097 \
+    --trust-remote-code --omni \
+    --stage-configs-path vllm_omni/deploy/audex_tts_30b.yaml
+# or: SIZE=30b MODE=tts examples/online_serving/audex/run_server.sh
+```
+
+Offline: same four scripts with
+`--model nvidia/Nemotron-Labs-Audex-30B-A3B --deploy-config vllm_omni/deploy/audex_<mode>_30b.yaml`.
+
+#### Verification
+
+Same curl / `client.py` as the 2B with
+`"model": "nvidia/Nemotron-Labs-Audex-30B-A3B"` (the server validates the
+model id). First launch downloads ~60 GB.
+
+#### Notes
+
+- Memory usage: ~61 GiB weights on one H100 80 GB (stage 0 at 0.85 GPU
+  fraction + decoder at 0.08 on the same card); verified with healthy KV
+  headroom. If long sequences OOM, set `tensor_parallel_size: 2` on
+  stage 0 (documented in the yaml headers).
+- `enable_prefix_caching: false` on stage 0 — hybrid Mamba constraint.
+- CUDA graphs stay ON for the thinker: the NemotronH stage captures
+  cleanly and eager mode costs ~7× single-stream TTS RTF (review-measured
+  0.30 vs 2.08).
+- Quality (en-24 self-transcribed CER): 0.72% unguided baseline; 0.72% at
+  cfg 1.5 with guided temperature 0.1.
+- Known limitation: the S2S chat pass needs the official
+  formatting-instruction prompt (the bundled examples/client send it);
+  a bare transcript turn makes the 30B answer in speech-codec tokens.

recipes/README.md

@@ -57,6 +57,7 @@ recipes/
 | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
 | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
 | [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
+| [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |

tests/config/test_omni_config.py

@@ -392,6 +392,11 @@ def test_sub_config_fields_match_rfc_scopes():
         "enable_multithread_weight_load",
         "num_weight_load_threads",
         "disable_autocast",
+        # Per-stage checkpoint resolution for repos whose stages live in
+        # subfolders (e.g. Audex): mirrors StagePipelineConfig on the
+        # legacy engine-args path.
+        "model_subdir",
+        "tokenizer_subdir",
     }
     assert {f.name for f in fields(OmniStageLoadConfig)} == {
         "load_format",

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -121,6 +121,10 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
             ("Qwen3TTSCode2Wav", "code2wav"),
             ("MiniCPMO45Code2Wav", "code2wav"),
             ("CosyVoice3Model", "cosyvoice3_code2wav"),
+            # audex: sync-path TTS decoder (async_chunk: false deploys) and
+            # the TTA XCodec1 stage (always sync full-payload).
+            ("AudexCode2Wav", "audex_code2wav"),
+            ("AudexXCodec1", "audex_xcodec"),
             ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
             ("DyninOmniForConditionalGeneration", "token2image"),
             ("DyninOmniForConditionalGeneration", "token2audio"),

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""HTTP + WebSocket validation for Qwen3-TTS Base and Higgs Audio V3 ``/v1/audio/speech`` (and related routes)."""
+"""HTTP + WebSocket validation for Qwen3-TTS Base, Higgs Audio V3 and Audex TTA ``/v1/audio/speech`` (and related routes)."""
 
 from __future__ import annotations
 
@@ -70,6 +70,19 @@ _HIGGS_AUDIO_V3_SPEECH = [
     ),
 ]
 
+_AUDEX_TTA_CAPTION = "Heavy rain falling on a tin roof."
+_AUDEX_TTA_SPEECH = [
+    pytest.param(
+        OmniServerParams(
+            model="nvidia/Nemotron-Labs-Audex-2B",
+            stage_config_path=get_deploy_config_path("audex_tta.yaml"),
+            server_args=_SPEECH_SERVER_ARGS,
+        ),
+        id="audex_tta",
+        marks=_L4_SPEECH_HW,
+    ),
+]
+
 
 def _apply_batch_overrides(body: dict[str, object], *, loc: str, overrides: dict[str, object]) -> None:
     if loc == "batch":
@@ -325,6 +338,37 @@ def test_speech_higgs_invalid_field_values(
     )
 
 
+# Audex TTA (``audex_tta`` deployment): the serving layer builds the caption
+# prompt and injects the RVQ phase contract + default guidance server-side, so
+# guidance knobs and speech-only fields must be rejected before generation.
+@pytest.mark.parametrize(
+    "overrides, err_message",
+    [
+        pytest.param({"extra_params": {"cfg_scale": 99}}, "cfg_scale", id="cfg_scale_out_of_range"),
+        pytest.param({"extra_params": {"cfg_scale": "big"}}, "cfg_scale", id="cfg_scale_wrong_type"),
+        pytest.param({"extra_params": {"tta_rvq": {}}}, "tta_rvq", id="tta_rvq_not_overridable"),
+        pytest.param({"input": "   "}, "caption", id="caption_whitespace_only"),
+        pytest.param({"voice": "alloy"}, "voice", id="voice_not_supported"),
+        pytest.param(
+            {"ref_audio": "https://example.com/ref.wav"},
+            "reference audio",
+            id="ref_audio_not_supported",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _AUDEX_TTA_SPEECH, indirect=True)
+def test_speech_audex_tta_invalid_field_values(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    body = {"model": omni_server.model, "input": _AUDEX_TTA_CAPTION, **overrides}
+    openai_client.send_audio_speech_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
 # ─────────────────────────────────────────────────────────────────────────────
 # POST /v1/audio/speech/batch (JSON)
 # ─────────────────────────────────────────────────────────────────────────────

tests/e2e/offline_inference/test_audex_s2s.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline E2E test for the Audex cascaded speech-to-speech pipeline.
+
+One ``audex_s2s`` deployment serves all three official
+passes. The routing contract is the hard gate here: text-modality passes
+(ASR, chat) finish at stage 0 and must never produce stage-1 audio; the
+audio-modality TTS pass streams through the causal speech decoder.
+"""
+
+from __future__ import annotations
+
+import copy
+import os
+import re
+import time
+
+import numpy as np
+import pytest
+from vllm.assets.audio import AudioAsset
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import concat_audio
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+MODEL_DIR_ENV = "VLLM_OMNI_AUDEX_MODEL_DIR"
+SAMPLE_RATE = 16_000
+
+ASR_PROMPT = (
+    "<|im_start|>user\n<so_embedding>\nTranscribe the input speech.<|im_end|>\n<|im_start|>assistant\n<think></think>"
+)
+
+# Official chat-pass recipe (examples/offline_inference/audex/speech_to_speech.py,
+# mirroring the official cascaded web demo) — 30B-A3B ONLY. The
+# formatting-instruction prefix is load-bearing on the 30B-A3B: without it a bare
+# transcript turn is treated as a TTS request and the model emits literal
+# <speechcodec_N> tokens instead of text, and under the yaml's near-greedy
+# defaults its reasoning block never closes. The 2B must NOT use this recipe:
+# it is long-verified stable on a bare prompt + engine-default sampling, while
+# the official temp-1.0 seed-0 draw hallucinates off-language text (yet closes
+# </think> cleanly, so the seed retry never fires) and fails the round-trip
+# WER gate — measured, reproduced byte-identically.
+SYSTEM_PROMPT = "You are a helpful and harmless assistant.\n\nYou are not allowed to use any tools."
+TEXT_INSTRUCTION = (
+    "SYSTEM & FORMATTING INSTRUCTION: You are Audex, created by NVIDIA based on the Nemotron-Cascade-2 architecture. "
+    "You may output your reasoning, followed by your final response. "
+    "You may format your reasoning block however you like. "
+    "However, your final response must strictly follow these rules: "
+    "* Write in plain, unformatted prose like a book or newspaper article. "
+    "* Do not use markdown, bullet points, lists, or headers. "
+    "* Standard numbers, abbreviations, and symbols are acceptable. "
+    "* [CRITICAL] You must press enter after every single sentence, placing each sentence on its own separate line."
+)
+CHAT_MAX_TOKENS = 2048  # official --text-max-tokens default
+# The official demo samples this pass unseeded; some trajectories deliberate
+# past the token budget without ever closing the think block, so retry a few
+# deterministic seeds until one closes.
+CHAT_SEEDS = (0, 1, 2)
+
+_audex_deployment = get_deploy_config_path("audex_s2s.yaml")
+_audex_model = os.environ.get(MODEL_DIR_ENV) or MODEL
+_OMNI_RUNNER_PARAMS = [
+    pytest.param(
+        (_audex_model, _audex_deployment, {"async_chunk": True}),
+        id="s2s_full",
+    ),
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    _OMNI_RUNNER_PARAMS.append(
+        pytest.param(
+            (_AUDEX_30B_MODEL, get_deploy_config_path("audex_s2s_30b.yaml"), {"async_chunk": True}),
+            id="s2s_30b",
+        )
+    )
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True),
+]
+
+
+def _normalize(text: str) -> list[str]:
+    return "".join(c.lower() if c.isalnum() or c.isspace() else " " for c in text).split()
+
+
+def _wer(hyp: str, ref: str) -> float:
+    hyp_words, ref_words = _normalize(hyp), _normalize(ref)
+    if not ref_words:
+        return 0.0
+    rows = list(range(len(hyp_words) + 1))
+    for i, ref_word in enumerate(ref_words, 1):
+        prev, rows[0] = rows[0], i
+        for j, hyp_word in enumerate(hyp_words, 1):
+            cur = min(rows[j] + 1, rows[j - 1] + 1, prev + (ref_word != hyp_word))
+            prev, rows[j] = rows[j], cur
+    return rows[-1] / len(ref_words)
+
+
+def _chat_prompt(user_text: str) -> str:
+    # Official chat-pass recipe: system prompt + formatting instruction
+    # prefixed to the user text, thinking ENABLED via an open <think> block
+    # (chat_template.jinja with enable_thinking=True renders
+    # "<|im_start|>assistant\n<think>\n"). Priming the block open keeps the
+    # reasoning inside <think>...</think> so _extract_final_response can drop
+    # it; with a closed block the 30B leaks the reasoning as plain text.
+    return (
+        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
+        f"<|im_start|>user\n{TEXT_INSTRUCTION}\n\n{user_text}<|im_end|>\n"
+        "<|im_start|>assistant\n<think>\n"
+    )
+
+
+def _extract_final_response(text: str) -> str:
+    # Mirror the official demo's extract_spoken_answer: keep only what
+    # follows the LAST closing think tag, falling back to the full text when
+    # the model skipped reasoning or never closed the block. Benign no-op for
+    # the 2B, whose answers carry no think block.
+    if "</think>" in text:
+        answer = text.rsplit("</think>", 1)[-1].strip()
+        if answer:
+            return answer
+    return text.strip()
+
+
+def _chat_params(runner: OmniRunner, seed: int):
+    # Official demo text-pass sampling (--text-temperature 1.0, --text-top-p
+    # 0.95, --text-max-tokens 2048). The yaml stage-0 defaults (temperature
+    # 0.1, top_k 80) are TTS-thinker settings; under near-greedy decoding the
+    # 30B loops inside its reasoning block until the token budget runs out
+    # and never emits </think>.
+    params = copy.deepcopy(runner.omni.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    stage0.temperature = 1.0
+    stage0.top_p = 0.95
+    stage0.top_k = -1
+    stage0.max_tokens = CHAT_MAX_TOKENS
+    stage0.seed = seed
+    return params
+
+
+def _chat_once(runner: OmniRunner, prompt: str, seed: int) -> tuple[str, bool]:
+    (req_output,) = runner.omni.generate([{"prompt": prompt, "modalities": ["text"]}], _chat_params(runner, seed))
+    out = req_output.outputs[0]
+    text = (out.text or "").strip()
+    truncated = getattr(out, "finish_reason", None) == "length" or len(out.token_ids or []) >= CHAT_MAX_TOKENS
+    return text, truncated
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_audex_s2s_cascade_round_trip(omni_runner: OmniRunner, run_level: str) -> None:
+    """ASR → chat → TTS over one deployment; text passes never emit audio."""
+    real_weights = run_level in {"advanced_model", "full_model"}
+    # Spoken question: vLLM's public audio asset (downloaded on first use).
+    audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
+    # Pass 1 — ASR (text-final).
+    asr_outputs = omni_runner.omni.generate(
+        [
+            {
+                "prompt": ASR_PROMPT,
+                "multi_modal_data": {"audio": (audio, sr)},
+                "modalities": ["text"],
+            }
+        ]
+    )
+    assert len(asr_outputs) == 1
+    asr_out = asr_outputs[0]
+    transcript = (asr_out.outputs[0].text or "").strip()
+    mm = getattr(asr_out, "multimodal_output", None) or {}
+    stage1_audio = concat_audio(mm.get("audio")) if "audio" in mm else np.zeros((0,), dtype=np.float32)
+    assert stage1_audio.size == 0, "Text-modality ASR pass produced stage-1 audio (routing regression)"
+
+    # Pass 2 — chat (text-final). The recipe is size-conditional (see the
+    # SYSTEM_PROMPT comment above): the 30B requires the official recipe,
+    # the 2B regresses under it.
+    if omni_runner.model_name == _AUDEX_30B_MODEL:
+        # Official recipe. The raw completion is
+        # <reasoning...></think><answer>; speak only the final response.
+        # Retry with the next seed only when the budget ran out before the
+        # reasoning block closed (a natural stop without </think> is a
+        # complete answer); if every seed spins, keep the full text
+        # (official demo behavior).
+        chat_prompt = _chat_prompt(transcript or "Hello.")
+        raw = ""
+        for seed in CHAT_SEEDS:
+            raw, truncated = _chat_once(omni_runner, chat_prompt, seed)
+            if "</think>" in raw or not truncated:
+                break
+        answer = _extract_final_response(raw)
+    else:
+        # 2B: bare prompt with a closed think block, engine-default sampling
+        # (yaml near-greedy), no retry; the answer is used as-is.
+        chat_prompt = f"<|im_start|>user\n{transcript or 'Hello.'}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+        (chat_output,) = omni_runner.omni.generate([{"prompt": chat_prompt, "modalities": ["text"]}])
+        answer = (chat_output.outputs[0].text or "").strip()
+
+    if real_weights:
+        assert transcript, "ASR pass produced an empty transcript"
+        assert answer, "Chat pass produced an empty answer"
+
+    # Pass 3 — TTS (audio-final; streams through code2wav). Both final
+    # stages may emit an output record; collect the audio-bearing one.
+    # The formatting instruction puts every sentence on its own line;
+    # flatten to plain prose for the TTS thinker.
+    tts_text = re.sub(r"\s+", " ", answer).strip() if answer else "The weather is nice today."
+    tts_text = tts_text[:200]
+    tts_prompt = {"prompt": build_cond_prompt(tts_text), "modalities": ["audio"]}
+    tts_outputs = omni_runner.omni.generate([tts_prompt])
+    assert tts_outputs, "TTS pass returned no outputs"
+    speech = np.zeros((0,), dtype=np.float32)
+    for req_output in tts_outputs:
+        mm = getattr(req_output, "multimodal_output", None) or {}
+        if "audio" in mm:
+            speech = concat_audio(mm["audio"])
+            break
+    assert speech.size > 0, "TTS pass produced empty audio"
+
+    if real_weights:
+        rms = float(np.sqrt(np.mean(np.square(speech))))
+        assert rms > 1e-3, f"TTS-pass audio near-silent (rms={rms})"
+
+        # Pass 4 — output verification: transcribe the synthesized answer
+        # with the SAME deployment's ASR pass and compare to the answer
+        # text (self-contained; no external ASR dependency).
+        verify_outputs = omni_runner.omni.generate(
+            [
+                {
+                    "prompt": ASR_PROMPT,
+                    "multi_modal_data": {"audio": (speech, SAMPLE_RATE)},
+                    "modalities": ["text"],
+                }
+            ]
+        )
+        verify_text = (verify_outputs[0].outputs[0].text or "").strip()
+        assert verify_text, "Output-verification ASR pass returned empty text"
+        # The full checkpoint prefixes transcripts with a language-tag
+        # sentence and quotes the content; extract the quoted payload.
+        if "'" in verify_text and len(verify_text.split("'")) >= 3:
+            verify_text = verify_text.split("'")[1].strip()
+        wer = _wer(verify_text, tts_text)
+        # Hard gate: the synthesized WAV must transcribe back to the answer
+        # (lenient threshold absorbs ASR quoting/prefix differences).
+        assert wer <= 0.5, f"Output WAV does not match the answer (WER={wer:.2f}): {verify_text!r} vs {tts_text!r}"
+        print(f"[trend] Audex S2S output-ASR WER vs answer: {wer:.3f} ({verify_text!r})")
+
+    # Final pass — time-to-first-audio for the streaming TTS path. This
+    # MUST be the last engine call: the py_generator API closes the engine
+    # when the generator is exhausted.
+    t_start = time.perf_counter()
+    first_audio_s: float | None = None
+    for req_output in omni_runner.omni.generate([tts_prompt], py_generator=True):
+        mm = getattr(req_output, "multimodal_output", None) or {}
+        if "audio" in mm and concat_audio(mm["audio"]).size > 0:
+            if first_audio_s is None:
+                first_audio_s = time.perf_counter() - t_start
+    assert first_audio_s is not None, "streaming TTS never produced an audio-bearing chunk"
+    print(f"[trend] Audex S2S TTS-pass first-audio latency: {first_audio_s * 1000.0:.0f} ms")

tests/e2e/offline_inference/test_audex_s2t.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline E2E smoke test for Audex audio understanding (thinker-only).
+
+Single-stage ``audex_thinker_only`` pipeline on
+``checkpoint_folder_full``: WAV (+ instruction) in, text out. The input clip
+is vLLM's public ``mary_had_lamb`` audio asset (downloaded on first use), so
+its transcript is known.
+
+WER is a trend metric (recorded in the assertion message, no numeric hard
+gate); the hard gates are structural: non-empty, non-degenerate text.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+from vllm.assets.audio import AudioAsset
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+MODEL_DIR_ENV = "VLLM_OMNI_AUDEX_MODEL_DIR"
+
+# Full transcript of the mary_had_lamb asset (Edison's phonograph recording);
+# a shortened reference would count the extra spoken words as insertions and
+# inflate WER past the trend threshold.
+_REFERENCE = (
+    "the first words i spoke in the original phonograph a little piece of practical poetry "
+    "mary had a little lamb its fleece was white as snow and everywhere that mary went the "
+    "lamb was sure to go"
+)
+
+ASR_PROMPT = (
+    "<|im_start|>user\n<so_embedding>\nTranscribe the input speech.<|im_end|>\n<|im_start|>assistant\n<think></think>"
+)
+
+_audex_deployment = get_deploy_config_path("audex_thinker_only.yaml")
+_audex_model = os.environ.get(MODEL_DIR_ENV) or MODEL
+_OMNI_RUNNER_PARAMS = [
+    pytest.param(
+        (_audex_model, _audex_deployment, {"async_chunk": False}),
+        id="thinker_only",
+    ),
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    _OMNI_RUNNER_PARAMS.append(
+        pytest.param(
+            (_AUDEX_30B_MODEL, get_deploy_config_path("audex_thinker_only_30b.yaml"), {"async_chunk": False}),
+            id="thinker_only_30b",
+        )
+    )
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True),
+]
+
+
+def _normalize(text: str) -> list[str]:
+    return "".join(c.lower() if c.isalnum() or c.isspace() else " " for c in text).split()
+
+
+def _wer(hyp: str, ref: str) -> float:
+    hyp_words, ref_words = _normalize(hyp), _normalize(ref)
+    if not ref_words:
+        return 0.0
+    rows = list(range(len(hyp_words) + 1))
+    for i, ref_word in enumerate(ref_words, 1):
+        prev, rows[0] = rows[0], i
+        for j, hyp_word in enumerate(hyp_words, 1):
+            cur = min(
+                rows[j] + 1,
+                rows[j - 1] + 1,
+                prev + (ref_word != hyp_word),
+            )
+            prev, rows[j] = rows[j], cur
+    return rows[-1] / len(ref_words)
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_audex_offline_asr_smoke(omni_runner: OmniRunner, run_level: str) -> None:
+    """Transcribing the public audio asset yields coherent, non-degenerate text."""
+    audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
+    outputs = omni_runner.omni.generate([{"prompt": ASR_PROMPT, "multi_modal_data": {"audio": (audio, sr)}}])
+
+    assert len(outputs) == 1
+    text = outputs[0].outputs[0].text if outputs[0].outputs else ""
+    assert isinstance(text, str)
+
+    if run_level in {"advanced_model", "full_model"}:
+        words = _normalize(text)
+        assert words, f"Empty transcription: {text!r}"
+        # Degenerate repetition guard: no word may dominate the output.
+        most_common = max(words.count(w) for w in set(words))
+        assert most_common <= max(3, len(words) // 2), f"Degenerate transcript: {text!r}"
+        wer = _wer(text, _REFERENCE)
+        # Trend metric (recorded, not gated): substantially-correct check only.
+        assert wer <= 0.9, f"Transcript unrelated to reference (WER={wer:.2f}): {text!r}"
+        print(f"[trend] Audex S2T WER vs reference: {wer:.3f} ({text!r})")

tests/e2e/offline_inference/test_audex_tta.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline E2E smoke test for Audex text-to-audio (TTA).
+
+Caption → RVQ-phase-masked <audiocodec_*> tokens → XCodec1 waveform through
+the ``audex_tta`` pipeline. RVQ phase validity is enforced by
+the mask logits processor (unit-tested against the official validator); the
+hard gates here are decode-side: non-silent, finite audio at the documented
+sample rate. CFG runs at the official TTA default (3.0).
+"""
+
+from __future__ import annotations
+
+import copy
+import os
+
+import numpy as np
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import concat_audio
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.models.audex.prompt import build_tta_cond_prompt
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+MODEL_DIR_ENV = "VLLM_OMNI_AUDEX_MODEL_DIR"
+SAMPLE_RATE = 16_000
+
+# The acceptance gate requires eight captions decoding to non-silent,
+# finite audio through XCodec1.
+CAPTIONS = (
+    "Heavy rain falling on a tin roof.",
+    "A dog barking in the distance while birds chirp.",
+    "Ocean waves crashing on a rocky shore.",
+    "A crackling campfire at night with crickets.",
+    "Thunder rumbling during a storm.",
+    "A steam train whistle blowing as it departs.",
+    "Wind chimes ringing in a gentle breeze.",
+    "Footsteps crunching through fresh snow.",
+)
+
+_audex_deployment = get_deploy_config_path("audex_tta.yaml")
+_audex_model = os.environ.get(MODEL_DIR_ENV) or MODEL
+_OMNI_RUNNER_PARAMS = [
+    pytest.param(
+        (_audex_model, _audex_deployment, {"async_chunk": False}),
+        id="tta_full_payload",
+    ),
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    _OMNI_RUNNER_PARAMS.append(
+        pytest.param(
+            (_AUDEX_30B_MODEL, get_deploy_config_path("audex_tta_30b.yaml"), {"async_chunk": False}),
+            id="tta_30b",
+        )
+    )
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True),
+]
+
+
+@pytest.fixture(scope="module")
+def tta_tokenizer_bits():
+    from transformers import AutoTokenizer
+
+    from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
+    from vllm_omni.model_executor.models.audex.tta import build_tta_phase_token_ids
+
+    root = ensure_audex_snapshot(_audex_model, profile="tta")
+    tokenizer = AutoTokenizer.from_pretrained(os.path.join(root, "checkpoint_folder_audiogen"), trust_remote_code=True)
+    phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer)
+    return tokenizer, phase_token_ids, start_tid, end_tid
+
+
+def _tta_sampling_params(runner: OmniRunner, cond_prompt: str, tokenizer_bits, cfg_scale: float = 3.0):
+    from vllm_omni.model_executor.models.audex.prompt import build_tta_null_prompt
+
+    tokenizer, phase_token_ids, start_tid, end_tid = tokenizer_bits
+    params = copy.deepcopy(runner.omni.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    if stage0.extra_args is None:
+        stage0.extra_args = {}
+    stage0.extra_args["tta_rvq"] = {
+        "phase_token_ids": phase_token_ids,
+        "start_tid": start_tid,
+        "end_tid": end_tid,
+        "codec_cap": 4000,
+        "start_in_prompt": True,
+    }
+    if cfg_scale > 1.0:
+        stage0.extra_args.update(
+            {
+                "cfg_scale": float(cfg_scale),
+                "cfg_role": "cond",
+                "cfg_pair_id": f"e2e-tta-{abs(hash(cond_prompt)) % 10_000_000}",
+                "cfg_null_prompt": build_tta_null_prompt(cond_prompt, tokenizer),
+            }
+        )
+    return params
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_audex_offline_tta_eight_captions(omni_runner: OmniRunner, run_level: str, tta_tokenizer_bits) -> None:
+    """All eight captions decode through XCodec1 to non-silent, finite audio.
+
+    RVQ phase validity is enforced by the mask logits processor by
+    construction; the final user output carries only decoded audio (the
+    audio-final pipeline emits no stage-0 token record), and disabling the
+    mask in-engine as an e2e control would require a second engine boot
+    with a different logits-processor config. The mask's sensitivity is
+    therefore proven at unit level (see test_audex_tta.py's mask tests and
+    the unmasked-sampling control) against the official validator.
+    """
+    real_weights = run_level in {"advanced_model", "full_model"}
+    for caption in CAPTIONS:
+        prompt = build_tta_cond_prompt(caption)
+        params = _tta_sampling_params(omni_runner, prompt, tta_tokenizer_bits)
+        outputs = omni_runner.omni.generate([prompt], params)
+
+        assert len(outputs) == 1, f"{caption!r}: expected one output, got {len(outputs)}"
+        audio_mm = outputs[0].multimodal_output
+        assert "audio" in audio_mm, f"{caption!r}: no audio output: {list(audio_mm.keys())}"
+        audio = concat_audio(audio_mm["audio"])
+        assert audio.size > 0, f"{caption!r}: generated audio is empty"
+        assert np.isfinite(audio).all(), f"{caption!r}: audio contains NaN/Inf"
+
+        sr_val = audio_mm.get("sr", SAMPLE_RATE)
+        if isinstance(sr_val, list) and sr_val:
+            sr_val = sr_val[-1]
+        if hasattr(sr_val, "item"):
+            sr_val = sr_val.item()
+        assert int(sr_val) == SAMPLE_RATE, f"{caption!r}: unexpected sample_rate={sr_val}"
+
+        if real_weights:
+            duration_s = audio.size / SAMPLE_RATE
+            # The decode cap is 500 frames = 10 s at 50 fps.
+            assert 0.1 <= duration_s <= 10.5, f"{caption!r}: unexpected duration={duration_s:.3f}s"
+            rms = float(np.sqrt(np.mean(np.square(audio))))
+            assert rms > 1e-4, f"{caption!r}: audio is near-silent (rms={rms})"

tests/e2e/offline_inference/test_audex_tts.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline E2E smoke test for Audex (Nemotron-Labs-Audex-2B) TTS.
+
+Passes the HF repo ROOT and verifies the 2-stage pipeline (thinker →
+streaming causal speech decoder): root-path/subfolder resolution, codec-token
+extraction, per-request streaming sessions, and the lookahead flush at EOS.
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import concat_audio
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+MODEL_DIR_ENV = "VLLM_OMNI_AUDEX_MODEL_DIR"
+SAMPLE_RATE = 16_000
+
+SYNTH_TEXTS = (
+    "The weather is so good, and I want to enjoy the beautiful morning in the park.",
+    "The quick brown fox jumps over the lazy dog.",
+)
+
+
+_audex_deployment = get_deploy_config_path("audex_tts.yaml")
+# Collection must not download anything: pass the repo id (or a local
+# override) and let the engine's stage-init path resolve/download the
+# required snapshot subset at execution time (ensure_audex_snapshot).
+_audex_model = os.environ.get(MODEL_DIR_ENV) or MODEL
+_OMNI_RUNNER_PARAMS = [
+    pytest.param(
+        (_audex_model, _audex_deployment, {"async_chunk": True}),
+        id="async_chunk",
+    ),
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    _OMNI_RUNNER_PARAMS.append(
+        pytest.param(
+            (_AUDEX_30B_MODEL, get_deploy_config_path("audex_tts_30b.yaml"), {"async_chunk": True}),
+            id="tts_30b",
+        )
+    )
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True),
+]
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_audex_offline_tts_smoke(omni_runner: OmniRunner, run_level: str) -> None:
+    """Audex TTS from the repo root should produce sane 16 kHz audio per prompt.
+
+    At ``core_model`` level the runner loads DUMMY weights (structural smoke
+    only); real-speech assertions apply from ``advanced_model`` up.
+    """
+    prompts = [build_cond_prompt(text) for text in SYNTH_TEXTS]
+    outputs = omni_runner.omni.generate(prompts)
+
+    real_weights = run_level in {"advanced_model", "full_model"}
+    assert len(outputs) == len(SYNTH_TEXTS), f"expected {len(SYNTH_TEXTS)} outputs, got {len(outputs)}"
+    for output in outputs:
+        audio_mm = output.multimodal_output
+        assert "audio" in audio_mm, f"No audio output found: {list(audio_mm.keys())}"
+        audio = concat_audio(audio_mm["audio"])
+        assert audio.size > 0, "Generated audio is empty"
+
+        sr_val = audio_mm.get("sr", SAMPLE_RATE)
+        if isinstance(sr_val, list) and sr_val:
+            sr_val = sr_val[-1]
+        if hasattr(sr_val, "item"):
+            sr_val = sr_val.item()
+        assert int(sr_val) == SAMPLE_RATE, f"Unexpected sample_rate={sr_val}"
+
+        if real_weights:
+            duration_s = audio.size / SAMPLE_RATE
+            assert 0.5 <= duration_s <= 20.0, f"Unexpected duration={duration_s:.3f}s"
+            # Real speech, not silence.
+            rms = float(np.sqrt(np.mean(np.square(audio))))
+            assert rms > 1e-3, f"Audio is near-silent (rms={rms})"
+
+
+def _cfg_sampling_params(runner: OmniRunner, cfg_scale: float, pair_id: str, cond_prompt: str):
+    """Stage sampling params carrying the CFG pair contract for one request."""
+    import copy
+
+    from transformers import AutoTokenizer
+
+    from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
+    from vllm_omni.model_executor.models.audex.prompt import build_null_prompt
+
+    # Resolve the ACTIVE param's model (2B or 30B), not the module default:
+    # the null prompt must be built with the same tokenizer the engine uses.
+    root = ensure_audex_snapshot(runner.model_name)
+    tokenizer = AutoTokenizer.from_pretrained(os.path.join(root, "checkpoint_folder_audiogen"), trust_remote_code=True)
+    params = copy.deepcopy(runner.omni.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    if stage0.extra_args is None:
+        stage0.extra_args = {}
+    stage0.extra_args.update(
+        {
+            "cfg_scale": float(cfg_scale),
+            "cfg_role": "cond",
+            "cfg_pair_id": pair_id,
+            "cfg_null_prompt": build_null_prompt(cond_prompt, tokenizer),
+        }
+    )
+    return params
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_audex_offline_cfg_guided_single_stream(omni_runner: OmniRunner, run_level: str) -> None:
+    """A guided request must yield exactly ONE audio stream (companion suppressed)."""
+    prompt = build_cond_prompt(SYNTH_TEXTS[0])
+    params = _cfg_sampling_params(omni_runner, 1.5, "e2e-cfg-pair", prompt)
+    outputs = omni_runner.omni.generate([prompt], params)
+
+    assert len(outputs) == 1, f"expected exactly one output (no companion leak), got {len(outputs)}"
+    audio = concat_audio(outputs[0].multimodal_output["audio"])
+    assert audio.size > 0, "Guided generation produced empty audio"
+    if run_level in {"advanced_model", "full_model"}:
+        rms = float(np.sqrt(np.mean(np.square(audio))))
+        assert rms > 1e-3, f"Guided audio is near-silent (rms={rms})"
+        duration_s = audio.size / SAMPLE_RATE
+        assert 0.5 <= duration_s <= 20.0, f"Unexpected duration={duration_s:.3f}s"
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_audex_offline_cfg_disabled_is_byte_identical(omni_runner: OmniRunner, run_level: str) -> None:
+    """cfg_scale=1.0 must be indistinguishable from not passing cfg at all."""
+    prompt = build_cond_prompt(SYNTH_TEXTS[1])
+
+    baseline = omni_runner.omni.generate([prompt])
+    base_audio = concat_audio(baseline[0].multimodal_output["audio"])
+
+    import copy
+
+    params = copy.deepcopy(omni_runner.omni.resolve_sampling_params_list(None))
+    if params[0].extra_args is None:
+        params[0].extra_args = {}
+    params[0].extra_args["cfg_scale"] = 1.0
+    disabled = omni_runner.omni.generate([prompt], params)
+    disabled_audio = concat_audio(disabled[0].multimodal_output["audio"])
+
+    assert base_audio.shape == disabled_audio.shape, (
+        f"cfg_scale=1.0 changed output length: {base_audio.shape} vs {disabled_audio.shape}"
+    )
+    np.testing.assert_array_equal(base_audio, disabled_audio)

tests/e2e/online_serving/test_audex_tta.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E online tests for Audex (Nemotron-Labs-Audex-2B) text-to-audio.
+
+Covers /v1/audio/speech through the ``audex_tta`` deployment:
+the ``audex_tta`` adapter builds the caption prompt and the serving layer
+injects the RVQ phase contract plus the official default guidance
+(cfg_scale 3.0) server-side.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+CAPTION = "Heavy rain falling on a tin roof."
+SAMPLE_RATE = 16_000
+
+audex_tta_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("audex_tta.yaml"),
+            server_args=["--trust-remote-code"],
+        ),
+        id="audex_tta",
+    )
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    audex_tta_server_params.append(
+        pytest.param(
+            OmniServerParams(
+                model=_AUDEX_30B_MODEL,
+                stage_config_path=get_deploy_config_path("audex_tta_30b.yaml"),
+                server_args=["--trust-remote-code"],
+            ),
+            id="online_tta_30b",
+        )
+    )
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_tta_server_params, indirect=True)
+def test_audex_tta_online_smoke(omni_server, openai_client, run_level) -> None:
+    """A caption returns a valid finite non-silent WAV at the default cfg 3.0."""
+    import io
+    import wave
+
+    import numpy as np
+    import requests
+
+    url = f"{openai_client.base_url.rstrip('/')}/v1/audio/speech"
+    resp = requests.post(
+        url,
+        json={"model": omni_server.model, "input": CAPTION, "response_format": "wav"},
+        timeout=600,
+    )
+    assert resp.status_code == 200, resp.text[:200]
+    assert resp.content[:4] == b"RIFF", "not a WAV response"
+
+    with wave.open(io.BytesIO(resp.content)) as wav:
+        assert wav.getframerate() == SAMPLE_RATE, f"unexpected sample rate {wav.getframerate()}"
+        frames = wav.readframes(wav.getnframes())
+        duration_s = wav.getnframes() / wav.getframerate()
+
+    audio = np.frombuffer(frames, dtype=np.int16).astype(np.float32) / 32768.0
+    assert audio.size > 0, "empty audio body"
+    assert np.isfinite(audio).all(), "audio contains NaN/Inf"
+
+    if run_level in {"advanced_model", "full_model"}:
+        # The decode cap is 500 frames = 10 s at 50 fps.
+        assert 0.1 <= duration_s <= 10.5, f"unexpected duration={duration_s:.3f}s"
+        rms = float(np.sqrt(np.mean(np.square(audio))))
+        assert rms > 1e-4, f"audio is near-silent (rms={rms})"

tests/e2e/online_serving/test_audex_tts.py

@@ -0,0 +1,230 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E online tests for Audex (Nemotron-Labs-Audex-2B) TTS.
+
+Covers /v1/audio/speech non-streaming and audio-byte streaming, plus the
+Audex request policies: single built-in voice (missing/empty/"default" only),
+no reference audio, and cfg_scale reserved at 1.0.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+PROMPT = "The weather is so good, and I want to enjoy the beautiful morning in the park."
+
+audex_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("audex_tts.yaml"),
+            server_args=["--trust-remote-code"],
+        ),
+        id="audex",
+    )
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    audex_server_params.append(
+        pytest.param(
+            OmniServerParams(
+                model=_AUDEX_30B_MODEL,
+                stage_config_path=get_deploy_config_path("audex_tts_30b.yaml"),
+                server_args=["--trust-remote-code"],
+            ),
+            id="online_tts_30b",
+        )
+    )
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_nonstream(omni_server, openai_client) -> None:
+    """Plain English TTS, non-streaming WAV response."""
+    request_config = {
+        "model": omni_server.model,
+        "input": PROMPT,
+        "stream": False,
+        "response_format": "wav",
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_stream_audio(omni_server, openai_client) -> None:
+    """Plain English TTS, raw audio-byte streaming."""
+    request_config = {
+        "model": omni_server.model,
+        "input": PROMPT,
+        "stream": True,
+        "stream_format": "audio",
+        "response_format": "wav",
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_max_new_tokens_caps_duration(omni_server, openai_client, run_level) -> None:
+    """max_new_tokens must cap stage-0 generation (review P2 regression)."""
+    import io
+    import wave
+
+    import requests
+
+    if run_level not in {"advanced_model", "full_model"}:
+        pytest.skip("duration cap needs real weights")
+
+    url = f"{openai_client.base_url.rstrip('/')}/v1/audio/speech"
+    resp = requests.post(
+        url,
+        json={
+            "model": omni_server.model,
+            "input": PROMPT,
+            "response_format": "wav",
+            "max_new_tokens": 25,
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    with wave.open(io.BytesIO(resp.content)) as wav:
+        duration_s = wav.getnframes() / wav.getframerate()
+    # 25 codec frames at 50 fps = 0.5 s (+ small lookahead tail); the full
+    # sentence would be ~3 s, so a capped result proves the override applied.
+    assert duration_s <= 1.0, f"max_new_tokens ignored: duration={duration_s:.2f}s"
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_request_policies(omni_server, openai_client) -> None:
+    """Unsupported voice / CFG / empty input must be rejected with 400."""
+    import requests
+
+    url = f"{openai_client.base_url.rstrip('/')}/v1/audio/speech"
+
+    ok = requests.post(url, json={"model": omni_server.model, "input": PROMPT, "voice": "default"})
+    assert ok.status_code == 200, ok.text
+
+    for bad_payload, needle in (
+        ({"input": PROMPT, "voice": "alloy"}, "voice"),
+        ({"input": "   "}, "non-empty"),
+        ({"input": PROMPT, "extra_params": {"cfg_scale": 0.5}}, "cfg_scale"),
+        ({"input": PROMPT, "extra_params": {"cfg_scale": 10.5}}, "cfg_scale"),
+        ({"input": PROMPT, "extra_params": {"cfg_scale": "big"}}, "cfg_scale"),
+        ({"input": PROMPT, "extra_params": {"cfg_role": "cond"}}, "cfg_role"),
+        ({"input": PROMPT, "extra_params": {"cfg_pair_id": "x"}}, "cfg_pair_id"),
+        ({"input": PROMPT, "ref_audio": "https://example.com/ref.wav"}, "reference audio"),
+    ):
+        resp = requests.post(url, json={"model": omni_server.model, **bad_payload})
+        assert resp.status_code == 400, f"{bad_payload} -> {resp.status_code}: {resp.text[:200]}"
+        assert needle in resp.text, f"{bad_payload} error missing {needle!r}: {resp.text[:200]}"
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_cfg_nonstream(omni_server, openai_client) -> None:
+    """Guided (cfg_scale=1.5) non-streaming request returns one WAV."""
+    request_config = {
+        "model": omni_server.model,
+        "input": PROMPT,
+        "stream": False,
+        "response_format": "wav",
+        "extra_params": {"cfg_scale": 1.5},
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_cfg_stream_audio(omni_server, openai_client) -> None:
+    """Guided (cfg_scale=1.5) streaming request yields a single audio stream."""
+    request_config = {
+        "model": omni_server.model,
+        "input": PROMPT,
+        "stream": True,
+        "stream_format": "audio",
+        "response_format": "wav",
+        "extra_params": {"cfg_scale": 1.5},
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_mixed_cfg_concurrency(omni_server, openai_client) -> None:
+    """Mixed guided/unguided requests at concurrency 4 all complete with audio."""
+    import concurrent.futures
+
+    import requests
+
+    url = f"{openai_client.base_url.rstrip('/')}/v1/audio/speech"
+    payloads = [
+        {"model": omni_server.model, "input": f"{PROMPT} Take {i}.", "response_format": "wav"}
+        | ({"extra_params": {"cfg_scale": 1.5}} if i % 2 == 0 else {})
+        for i in range(4)
+    ]
+
+    def _post(payload):
+        return requests.post(url, json=payload, timeout=300)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
+        responses = list(pool.map(_post, payloads))
+
+    for payload, resp in zip(payloads, responses):
+        assert resp.status_code == 200, f"{payload} -> {resp.status_code}: {resp.text[:200]}"
+        assert resp.content[:4] == b"RIFF", f"{payload}: not a WAV response"
+        assert len(resp.content) > 44, f"{payload}: empty audio body"
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_stream_abort_leaves_server_healthy(omni_server, openai_client) -> None:
+    """Aborting a guided streaming request mid-flight must not wedge the server.
+
+    The abort cascades to the CFG companion (cfg_companion_tracker); a
+    follow-up guided request on the same server must still succeed.
+    """
+    import requests
+
+    url = f"{openai_client.base_url.rstrip('/')}/v1/audio/speech"
+    long_input = " ".join([PROMPT] * 4)
+
+    with requests.post(
+        url,
+        json={
+            "model": omni_server.model,
+            "input": long_input,
+            "stream": True,
+            "stream_format": "audio",
+            "response_format": "wav",
+            "extra_params": {"cfg_scale": 1.5},
+        },
+        stream=True,
+        timeout=300,
+    ) as resp:
+        assert resp.status_code == 200, resp.text
+        # Read one chunk, then abort the connection mid-stream.
+        next(resp.iter_content(chunk_size=4096))
+
+    follow_up = requests.post(
+        url,
+        json={
+            "model": omni_server.model,
+            "input": PROMPT,
+            "response_format": "wav",
+            "extra_params": {"cfg_scale": 1.5},
+        },
+        timeout=300,
+    )
+    assert follow_up.status_code == 200, follow_up.text[:200]
+    assert follow_up.content[:4] == b"RIFF" and len(follow_up.content) > 44

tests/entrypoints/openai_api/test_audex_serving_guards.py

@@ -0,0 +1,376 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Serving-layer regressions for the Audex zero-codec-token contract.
+
+A zero-codec-token (or phase-invalid TTA) request reaches the serving layer
+as a stream whose only audio payloads are zero-size tensors. The streaming
+generator must not turn that into a header-only/empty successful stream: the
+raw-bytes generator raises (before any WAV header), and the SSE wrapper
+surfaces ``speech.audio.error`` instead of ``speech.audio.done``. The
+contract covers both audex pipelines (``audex`` TTS and ``audex_tta``).
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.entrypoints.openai.serving_speech import (
+    _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES,
+    OmniOpenAIServingSpeech,
+)
+
+
+def _serving(model_type: str = "audex") -> OmniOpenAIServingSpeech:
+    serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
+    serving._tts_model_type = model_type
+    serving._mark_ref_audio_artifact_ready_for_request = lambda request_id: None
+    serving._discard_ref_audio_artifact_warmup = lambda request_id: None
+    return serving
+
+
+def _result(audio: torch.Tensor) -> SimpleNamespace:
+    return SimpleNamespace(multimodal_output={"model_outputs": audio, "sr": torch.tensor(16000, dtype=torch.int32)})
+
+
+async def _gen(*results):
+    for res in results:
+        yield res
+
+
+async def _collect(chunks):
+    out = []
+    async for chunk in chunks:
+        out.append(chunk)
+    return out
+
+
+def test_no_audio_guard_covers_both_audex_pipelines():
+    # The non-streaming CreateAudio guard keys off the same constant; pin
+    # its membership so neither pipeline regresses to empty-200 responses.
+    assert _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES == {"audex", "audex_tta"}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_type", ["audex", "audex_tta"])
+async def test_raw_stream_raises_on_empty_audio_before_any_header(model_type):
+    serving = _serving(model_type)
+    chunks = serving._generate_audio_chunks(
+        _gen(_result(torch.empty(0)), _result(torch.empty(0))),
+        request_id="req-empty",
+        response_format="wav",
+    )
+    with pytest.raises(ValueError, match="no audio"):
+        await _collect(chunks)
+
+
+@pytest.mark.asyncio
+async def test_raw_stream_yields_audio_normally_when_non_empty():
+    serving = _serving()
+    chunks = serving._generate_audio_chunks(
+        _gen(_result(torch.zeros(320)), _result(torch.empty(0))),
+        request_id="req-ok",
+        response_format="wav",
+    )
+    collected = await _collect(chunks)
+    # WAV header + one PCM chunk; the trailing empty chunk is skipped silently.
+    assert len(collected) == 2
+    assert collected[0][:4] == b"RIFF"
+    assert len(collected[1]) == 320 * 2  # int16 PCM
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_type", ["audex", "audex_tta"])
+async def test_sse_stream_emits_error_not_done_on_empty_audio(model_type):
+    serving = _serving(model_type)
+    events = serving._generate_audio_sse_events(
+        _gen(_result(torch.empty(0))),
+        request_id="req-empty",
+        response_format="pcm",
+    )
+    collected = await _collect(events)
+    assert any("speech.audio.error" in event for event in collected)
+    assert not any("speech.audio.done" in event for event in collected)
+
+
+@pytest.mark.asyncio
+async def test_non_audex_empty_chunks_keep_legacy_behavior():
+    """Other models keep their existing semantics (no new rejection)."""
+    serving = _serving(model_type="qwen3_tts")
+    chunks = serving._generate_audio_chunks(
+        _gen(_result(torch.empty(0))),
+        request_id="req-other",
+        response_format="pcm",
+    )
+    await _collect(chunks)  # must not raise
+
+
+# ---------------------------------------------------------------- CFG serving surface
+
+import re  # noqa: E402
+
+from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest  # noqa: E402
+from vllm_omni.entrypoints.openai.tts_adapters.audex import AudexAdapter  # noqa: E402
+
+
+def _adapter() -> AudexAdapter:
+    return AudexAdapter.__new__(AudexAdapter)
+
+
+def _speech_request(**extra_params) -> OpenAICreateSpeechRequest:
+    payload = {"model": "audex", "input": "Hello world."}
+    if extra_params:
+        payload["extra_params"] = extra_params
+    return OpenAICreateSpeechRequest(**payload)
+
+
+class _MarkerTokenizer:
+    """Special markers count one token; words and punctuation split."""
+
+    _TOKEN_RE = re.compile(r"<unk>|<\|[^|]*\|>|<[a-z_]+>|\w+|[^\w\s]")
+
+    def encode(self, text: str) -> list[int]:
+        return [0] * len(self._TOKEN_RE.findall(text))
+
+
+class TestAudexCfgValidation:
+    def test_no_cfg_scale_accepted(self):
+        assert _adapter().validate(_speech_request()) is None
+
+    @pytest.mark.parametrize("scale", [1.0, 1.5, 3.0, 10.0])
+    def test_in_range_scale_accepted(self, scale):
+        assert _adapter().validate(_speech_request(cfg_scale=scale)) is None
+
+    @pytest.mark.parametrize("scale", [0.5, 0.0, -1.0, 10.5, "big"])
+    def test_out_of_range_or_malformed_scale_rejected(self, scale):
+        error = _adapter().validate(_speech_request(cfg_scale=scale))
+        assert error is not None and "cfg_scale" in error
+
+    def test_explicit_null_scale_treated_as_absent(self):
+        assert _adapter().validate(_speech_request(cfg_scale=None)) is None
+        serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
+        params = SimpleNamespace(extra_args={"cfg_scale": None})
+        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params)
+        assert out.extra_args == {}
+
+    @pytest.mark.parametrize("key", ["cfg_role", "cfg_pair_id", "cfg_null_prompt"])
+    def test_internal_cfg_keys_rejected(self, key):
+        error = _adapter().validate(_speech_request(**{key: "x"}))
+        assert error is not None and key in error
+
+
+class TestAudexCfgInjection:
+    def _serving_with_fake_tokenizer(self):
+        serving = _serving()
+        serving._get_audex_tokenizer = lambda: _MarkerTokenizer()
+        return serving
+
+    def _cond_prompt(self, text: str = "Hello world.") -> str:
+        from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+        return build_cond_prompt(text)
+
+    def test_no_extra_args_is_untouched(self):
+        serving = self._serving_with_fake_tokenizer()
+        params = SimpleNamespace(extra_args=None)
+        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
+        assert out.extra_args is None
+        assert params.extra_args is None
+
+    def test_scale_one_leaves_params_byte_identical(self):
+        serving = self._serving_with_fake_tokenizer()
+        params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
+        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
+        assert out.extra_args == {}
+
+    def test_guided_request_gains_pair_contract(self):
+        serving = self._serving_with_fake_tokenizer()
+        params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
+        cond = self._cond_prompt("Some transcription text here.")
+        out = serving._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params)
+
+        extra = out.extra_args
+        assert extra["cfg_scale"] == 1.5
+        assert extra["cfg_role"] == "cond"
+        assert extra["cfg_pair_id"] == "req-42"
+        null_prompt = extra["cfg_null_prompt"]
+        assert "<unk>" in null_prompt
+        assert "Some transcription" not in null_prompt
+        tok = _MarkerTokenizer()
+        assert len(tok.encode(null_prompt)) == len(tok.encode(cond))
+
+    def test_missing_text_prompt_raises(self):
+        serving = self._serving_with_fake_tokenizer()
+        params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
+        with pytest.raises(ValueError, match="adapter-built text prompt"):
+            serving._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params)
+
+    def test_injection_never_mutates_input_params(self):
+        """The input may be the engine's SHARED default params (review P2):
+        pair state must land only on the returned clone."""
+        serving = self._serving_with_fake_tokenizer()
+        shared = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
+        out = serving._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared)
+        assert out is not shared
+        assert shared.extra_args == {"cfg_scale": 1.5}
+        assert shared.temperature == 0.1
+        assert out.extra_args["cfg_pair_id"] == "req-7"
+        assert out.temperature == 0.05
+
+
+class TestAudexStageDetection:
+    """audex_omni is speech-capable only alongside the speech decoder."""
+
+    def _serving_with_stages(self, stages: list[str]) -> OmniOpenAIServingSpeech:
+        serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
+        serving.engine_client = SimpleNamespace(
+            stage_configs=[SimpleNamespace(engine_args=SimpleNamespace(model_stage=s)) for s in stages]
+        )
+        return serving
+
+    def test_full_pipeline_detected_as_tts(self):
+        serving = self._serving_with_stages(["audex_omni", "audex_code2wav"])
+        stage = serving._find_tts_stage()
+        assert stage is not None and stage.engine_args.model_stage == "audex_omni"
+
+    def test_thinker_only_not_detected_as_tts(self):
+        serving = self._serving_with_stages(["audex_omni"])
+        assert serving._find_tts_stage() is None
+
+    def test_tts_pipeline_still_detected(self):
+        serving = self._serving_with_stages(["audex_thinker", "audex_code2wav"])
+        stage = serving._find_tts_stage()
+        assert stage is not None and stage.engine_args.model_stage == "audex_thinker"
+
+
+def test_guided_injection_sets_guided_temperature():
+    """Guided requests decode at the measured CFG temperature (0.05)."""
+    serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
+    cond = TestAudexCfgInjection()._cond_prompt()
+    params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
+    out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params)
+    assert out.temperature == 0.05
+
+    unguided = SimpleNamespace(extra_args={"cfg_scale": 1.0}, temperature=0.1)
+    out = serving._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided)
+    assert out.temperature == 0.1
+
+
+# ---------------------------------------------------------------- TTA serving surface
+
+from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter  # noqa: E402
+from vllm_omni.entrypoints.openai.tts_adapters.audex_tta import AudexTTAAdapter  # noqa: E402
+
+
+def _tta_adapter() -> AudexTTAAdapter:
+    return AudexTTAAdapter.__new__(AudexTTAAdapter)
+
+
+class _TTAVocabTokenizer(_MarkerTokenizer):
+    """Marker tokenizer + a full audiocodec vocab for phase-id building."""
+
+    def get_vocab(self):
+        vocab = {"<audiogen_start>": 7, "<audiogen_end>": 8}
+        for n in range(8192):
+            vocab[f"<audiocodec_{n}>"] = 1000 + n
+        return vocab
+
+
+class TestAudexTTAAdapterValidation:
+    def test_default_caption_accepted(self):
+        assert _tta_adapter().validate(_speech_request()) is None
+
+    def test_empty_caption_rejected(self):
+        req = _speech_request()
+        req.input = "  "
+        assert _tta_adapter().validate(req) is not None
+
+    @pytest.mark.parametrize("scale", [1.0, 3.0, 10.0])
+    def test_in_range_scale_accepted(self, scale):
+        assert _tta_adapter().validate(_speech_request(cfg_scale=scale)) is None
+
+    @pytest.mark.parametrize("scale", [0.5, 10.5, "big"])
+    def test_out_of_range_scale_rejected(self, scale):
+        error = _tta_adapter().validate(_speech_request(cfg_scale=scale))
+        assert error is not None and "cfg_scale" in error
+
+    @pytest.mark.parametrize("key", ["cfg_role", "cfg_pair_id", "cfg_null_prompt", "tta_rvq"])
+    def test_internal_keys_rejected(self, key):
+        error = _tta_adapter().validate(_speech_request(**{key: "x"}))
+        assert error is not None and key in error
+
+
+class TestAdapterRouting:
+    def test_tta_adapter_resolves_by_model_type(self):
+        adapter_cls = resolve_adapter("audex_tta")
+        assert adapter_cls is AudexTTAAdapter
+        assert adapter_cls.stage_keys == frozenset({"audex_tta_thinker"})
+
+    def test_tts_adapter_never_selects_tta_stage(self):
+        assert "audex_tta_thinker" not in AudexAdapter.stage_keys
+        assert not (AudexAdapter.stage_keys & AudexTTAAdapter.stage_keys)
+
+    def test_detection_separates_tts_and_tta(self):
+        tta_serving = TestAudexStageDetection()._serving_with_stages(["audex_tta_thinker", "audex_xcodec"])
+        tta_serving._tts_stage = tta_serving._find_tts_stage()
+        assert tta_serving._tts_stage.engine_args.model_stage == "audex_tta_thinker"
+        assert tta_serving._detect_tts_model_type() == "audex_tta"
+
+        tts_serving = TestAudexStageDetection()._serving_with_stages(["audex_thinker", "audex_code2wav"])
+        tts_serving._tts_stage = tts_serving._find_tts_stage()
+        assert tts_serving._detect_tts_model_type() == "audex"
+
+
+class TestAudexTTAInjection:
+    def _serving(self):
+        serving = _serving(model_type="audex_tta")
+        serving._get_audex_tokenizer = lambda: _TTAVocabTokenizer()
+        serving._audex_tta_rvq = None
+        return serving
+
+    def _prompt(self):
+        from vllm_omni.model_executor.models.audex.prompt import build_tta_cond_prompt
+
+        return {"prompt": build_tta_cond_prompt("Rain on a tin roof.")}
+
+    def test_default_injection_applies_rvq_and_cfg3(self):
+        serving = self._serving()
+        params = SimpleNamespace(extra_args=None)
+        out = serving._inject_audex_tta_args("req-9", self._prompt(), params)
+
+        extra = out.extra_args
+        rvq = extra["tta_rvq"]
+        assert len(rvq["phase_token_ids"]) == 4
+        assert all(len(ids) == 1024 for ids in rvq["phase_token_ids"])
+        assert (rvq["start_tid"], rvq["end_tid"]) == (7, 8)
+        assert rvq["codec_cap"] == 4000 and rvq["start_in_prompt"] is True
+        assert extra["cfg_scale"] == 3.0
+        assert extra["cfg_role"] == "cond" and extra["cfg_pair_id"] == "req-9"
+        assert "<unk>" in extra["cfg_null_prompt"]
+
+    def test_explicit_scale_one_disables_cfg_but_keeps_rvq(self):
+        serving = self._serving()
+        params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
+        out = serving._inject_audex_tta_args("req-1", self._prompt(), params)
+        assert "tta_rvq" in out.extra_args
+        assert "cfg_scale" not in out.extra_args
+        assert "cfg_role" not in out.extra_args
+
+    def test_rvq_contract_cached_per_process(self):
+        serving = self._serving()
+        out_a = serving._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None))
+        out_b = serving._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None))
+        assert out_a.extra_args["tta_rvq"] is out_b.extra_args["tta_rvq"]
+
+    def test_injection_never_mutates_input_params(self):
+        """The default online path hands the injector an element of the
+        engine's SHARED default_sampling_params_list (review P2): the
+        per-request pair state must land only on the returned clone."""
+        serving = self._serving()
+        shared = SimpleNamespace(extra_args=None)
+        out_1 = serving._inject_audex_tta_args("req-1", self._prompt(), shared)
+        out_2 = serving._inject_audex_tta_args("req-2", self._prompt(), shared)
+        assert shared.extra_args is None
+        assert out_1 is not shared and out_2 is not shared
+        assert out_1.extra_args["cfg_pair_id"] == "req-1"
+        assert out_2.extra_args["cfg_pair_id"] == "req-2"

tests/helpers/media.py

@@ -529,6 +529,26 @@ def decode_b64_image(b64: str):
     return img
 
 
+def concat_audio(audio_val) -> np.ndarray:
+    """Flatten a multimodal audio payload to mono float32 samples.
+
+    Engines return ``multimodal_output["audio"]`` as a tensor, a list of
+    per-chunk tensors (streaming decoders), or an array-like; concatenate
+    in order and return a 1-D ``np.float32`` array (empty when a list has
+    no tensors).
+    """
+    import torch
+
+    if isinstance(audio_val, list):
+        tensors = [t.detach().cpu().float().reshape(-1) for t in audio_val if isinstance(t, torch.Tensor)]
+        if not tensors:
+            return np.zeros((0,), dtype=np.float32)
+        return torch.cat(tensors, dim=-1).numpy().astype(np.float32, copy=False)
+    if isinstance(audio_val, torch.Tensor):
+        return audio_val.detach().cpu().float().reshape(-1).numpy()
+    return np.asarray(audio_val, dtype=np.float32).reshape(-1)
+
+
 def preprocess_text(text):
     import opencc
 
@@ -695,6 +715,7 @@ def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small") ->
 
 __all__ = [
     "_merge_base64_audio_to_segment",
+    "concat_audio",
     "convert_audio_bytes_to_text",
     "convert_audio_file_to_text",
     "cosine_similarity_text",

vllm_omni/config/config_factory.py

@@ -45,6 +45,26 @@ logger = init_logger(__name__)
 _DEFAULT_PARALLEL_DEGREE = 1
 
 
+def with_trust_remote_code_override(
+    overrides: Mapping[str, Any],
+    trust_remote_code: bool | None,
+) -> dict[str, Any]:
+    """Merge the tri-state ``trust_remote_code`` into an override mapping.
+
+    Single home for the precedence rule (explicit caller value > deploy
+    yaml per-stage value > vLLM default False): a non-None value becomes an
+    explicit override; ``None`` means "not specified" and leaves the deploy
+    yaml's per-stage setting in effect. The serve ``--trust-remote-code``
+    flag is store_true — its absent-False must be mapped to ``None`` at the
+    CLI boundary before reaching here, since it cannot express an explicit
+    False.
+    """
+    merged = dict(overrides)
+    if trust_remote_code is not None:
+        merged["trust_remote_code"] = trust_remote_code
+    return merged
+
+
 class StageConfigFactory:
     """Factory that loads pipeline YAML and merges CLI overrides.
 
@@ -293,7 +313,7 @@ class StageConfigFactory:
         cls,
         model: str,
         *,
-        trust_remote_code: bool,
+        trust_remote_code: bool | None,
         cli_overrides: dict[str, Any],
         deploy_config_path: str | None,
     ) -> VllmOmniConfig | None:
@@ -301,18 +321,19 @@ class StageConfigFactory:
         user_deploy_config = cls._load_user_deploy_config(deploy_config_path)
         pipeline_cfg = cls.get_pipeline_config(
             model=model,
-            trust_remote_code=trust_remote_code,
+            # HF config resolution needs a real bool: transformers treats
+            # None as "prompt for consent", which blocks non-interactively.
+            trust_remote_code=bool(trust_remote_code),
             deploy_config_path=deploy_config_path,
             user_deploy_config=user_deploy_config,
         )
         if pipeline_cfg is None:
             return None
 
-        registry_cli_overrides = {
-            **cli_overrides,
-            "trust_remote_code": trust_remote_code,
-            "model": model,
-        }
+        registry_cli_overrides = with_trust_remote_code_override(
+            {**cli_overrides, "model": model},
+            trust_remote_code,
+        )
         return VllmOmniConfig.from_pipeline_config(
             pipeline_cfg,
             user_deploy_config=user_deploy_config,
@@ -325,7 +346,7 @@ class StageConfigFactory:
         cls,
         model: str,
         *,
-        trust_remote_code: bool,
+        trust_remote_code: bool | None,
         cli_overrides: dict[str, Any],
         deploy_config_path: str | None,
         strategy_specs: Mapping[Any, Any] | None = None,
@@ -339,17 +360,15 @@ class StageConfigFactory:
         user_deploy_config = cls._load_user_deploy_config(deploy_config_path)
         pipeline_cfg = cls.get_pipeline_config(
             model=model,
-            trust_remote_code=trust_remote_code,
+            # See create_from_model: HF resolution needs a real bool.
+            trust_remote_code=bool(trust_remote_code),
             deploy_config_path=deploy_config_path,
             user_deploy_config=user_deploy_config,
         )
         if pipeline_cfg is None:
             return None, None
 
-        legacy_cli_overrides = {
-            **cli_overrides,
-            "trust_remote_code": trust_remote_code,
-        }
+        legacy_cli_overrides = with_trust_remote_code_override(cli_overrides, trust_remote_code)
         return cls._create_legacy_from_registry(
             pipeline_cfg,
             legacy_cli_overrides,

vllm_omni/config/model.py

@@ -146,7 +146,10 @@ class OmniModelConfig(ModelConfig):
 
     @property
     def architectures(self) -> list[str]:
-        if self.model_arch is not None:
+        # Falsy (None or "") means "no stage override": fall back to the
+        # checkpoint config's own architectures. The stage-config builder
+        # emits None; "" is tolerated for legacy callers.
+        if self.model_arch:
             return [self.model_arch]
         return super().architectures
 

vllm_omni/config/omni_config.py

@@ -270,6 +270,11 @@ class OmniStageModelConfig:
     enable_multithread_weight_load: bool = True
     num_weight_load_threads: int = Field(default=4, ge=1)
     disable_autocast: bool = False
+    # Per-stage checkpoint/tokenizer subdirectories under the model root
+    # (e.g. Audex stage 0 → checkpoint_folder_audiogen). Mirrors
+    # StagePipelineConfig.model_subdir/tokenizer_subdir on the legacy path.
+    model_subdir: str | None = None
+    tokenizer_subdir: str | None = None
 
 
 @config
@@ -1117,6 +1122,10 @@ def _build_model_config(
     kwargs = _config_kwargs(engine)
     if "has_sampling_extra_args" not in kwargs:
         kwargs["has_sampling_extra_args"] = bool((default_sampling_params or {}).get("extra_args"))
+    if "model_subdir" not in kwargs and topology.model_subdir is not None:
+        kwargs["model_subdir"] = topology.model_subdir
+    if "tokenizer_subdir" not in kwargs and topology.tokenizer_subdir is not None:
+        kwargs["tokenizer_subdir"] = topology.tokenizer_subdir
     return OmniStageModelConfig(
         default_sampling_params=default_sampling_params,
         **kwargs,

vllm_omni/config/pipeline_registry.py

@@ -33,6 +33,12 @@ from vllm.logger import init_logger
 from vllm_omni.config.stage_config import (
     PipelineConfig,
 )
+from vllm_omni.model_executor.models.audex.pipeline import (
+    AUDEX_S2S_PIPELINE,
+    AUDEX_THINKER_ONLY_PIPELINE,
+    AUDEX_TTA_PIPELINE,
+    AUDEX_TTS_PIPELINE,
+)
 from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE
 from vllm_omni.model_executor.models.bagel.pipeline import (
     BAGEL_PIPELINE,
@@ -126,6 +132,14 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE,
     "voxcpm2": VOXCPM2_PIPELINE,
     "cosyvoice3": COSYVOICE3_PIPELINE,
+    "audex_tts": AUDEX_TTS_PIPELINE,
+    "audex_tta": AUDEX_TTA_PIPELINE,
+    "audex_thinker_only": AUDEX_THINKER_ONLY_PIPELINE,
+    "audex_s2s": AUDEX_S2S_PIPELINE,
+    # Alias: the Nemotron-Labs-Audex-2B repo-root config.json reports
+    # ``model_type: nemotron_labs_audex``; bare ``vllm-omni serve <repo>``
+    # auto-detects through it and must land on the default (TTS) pipeline.
+    "nemotron_labs_audex": AUDEX_TTS_PIPELINE,
     "mimo_audio": MIMO_AUDIO_PIPELINE,
     "ming_tts": MING_TTS_PIPELINE,
     "ming_tts_moe": MING_TTS_MOE_PIPELINE,

vllm_omni/config/stage_config.py

@@ -821,7 +821,7 @@ def _build_engine_args(
     per-stage StageDeployConfig overrides take precedence when present (e.g.
     ``engine_extras`` can still carry a stage-specific ``dtype``).
     """
-    engine_args: dict[str, Any] = {"model_arch": ps.model_arch or pipeline.model_arch}
+    engine_args: dict[str, Any] = {"model_arch": ps.model_arch or pipeline.model_arch or None}
     if ps.execution_type == StageExecutionType.DIFFUSION and ps.model_arch:
         engine_args.setdefault("model_class_name", ps.model_arch)
     if ps.engine_output_type:

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -51,6 +51,12 @@ _FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
         ("MiniCPMO45Code2Wav", "code2wav"),
         # cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1).
         ("CosyVoice3Model", "cosyvoice3_code2wav"),
+        # audex TTS sync path: thinker (Stage 0) -> streaming decoder (Stage 1).
+        # The default deploy is async_chunk; this covers async_chunk: false.
+        ("AudexCode2Wav", "audex_code2wav"),
+        # audex TTA: tta thinker (Stage 0) -> XCodec1 (Stage 1, always sync
+        # full-payload — CNN codec decoded over the full sequence).
+        ("AudexXCodec1", "audex_xcodec"),
         # indextts2: indextts2_talker (Stage 0) -> indextts2_s2mel_decoder
         # (Stage 1). Stage 1 consumes the complete mel/latent payload.
         ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),

vllm_omni/deploy/audex_s2s.yaml

@@ -0,0 +1,65 @@
+# Audex (Nemotron-Labs-Audex-2B) full deploy: cascaded speech-to-speech.
+#
+# One audio-capable thinker (checkpoint_folder_full) serves the official
+# three-pass cascade — (1) ASR: audio -> transcript, (2) chat: transcript ->
+# answer, (3) TTS: answer -> speech — plus the streaming causal speech
+# decoder as stage 1. Per-request output modalities route the passes: text
+# passes finish at stage 0; only audio-modality requests stream codec frames
+# into code2wav. CFG applies to the TTS pass via extra_params.cfg_scale
+# (recommended 1.5; 1.0 = off).
+pipeline: audex_s2s
+async_chunk: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      codec_chunk_frames: 25
+      initial_codec_chunk_frames: 1
+      codec_token_offset: 131077
+      codec_vocab_size: 65536
+      decoder_emit_chunk_frames: 5
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.5
+    trust_remote_code: true
+    enable_prefix_caching: true
+    async_scheduling: true
+    max_num_batched_tokens: 8192
+    # Sound embeddings (up to 30 clips x 750) + speech-token generation.
+    max_model_len: 32768
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.1
+      top_k: 80
+      max_tokens: 2048
+      seed: 0
+
+  - stage_id: 1
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/deploy/audex_s2s_30b.yaml

@@ -0,0 +1,60 @@
+# Audex 30B-A3B full deploy: cascaded speech-to-speech.
+# See audex_s2s.yaml for the cascade/routing notes and audex_tts_30b.yaml
+# for the 30B sizing notes (hybrid Mamba: no prefix caching; single-H100
+# layout, TP2 fallback; pass this yaml explicitly).
+pipeline: audex_s2s
+async_chunk: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      codec_chunk_frames: 25
+      initial_codec_chunk_frames: 1
+      codec_token_offset: 131077
+      codec_vocab_size: 65536
+      decoder_emit_chunk_frames: 5
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.85
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 8192
+    # Sound embeddings (up to 30 clips x 750) + speech-token generation.
+    max_model_len: 32768
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.1
+      top_k: 80
+      max_tokens: 2048
+      seed: 0
+
+  - stage_id: 1
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.08
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/deploy/audex_thinker_only.yaml

@@ -0,0 +1,25 @@
+# Audex (Nemotron-Labs-Audex-2B) thinker-only deploy: audio understanding
+# (ASR / audio QA) -> text. Single-stage configuration of the audio-capable
+# full checkpoint (checkpoint_folder_full): NV-Whisper encoder + relu2
+# projector + 2B dense LM.
+#
+# The official audioqa script decodes ASR at temperature 0; max_tokens 256
+# covers transcription and short QA answers (raise for long-form answers).
+pipeline: audex_thinker_only
+async_chunk: false
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.7
+    trust_remote_code: true
+    enable_prefix_caching: true
+    max_num_batched_tokens: 8192
+    # 30 clips x 750 sound embeddings (900 s cap) + text headroom.
+    max_model_len: 32768
+    dtype: bfloat16
+    devices: "0"
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 256
+      seed: 0

vllm_omni/deploy/audex_thinker_only_30b.yaml

@@ -0,0 +1,22 @@
+# Audex 30B-A3B thinker-only deploy: audio understanding (ASR / audio QA).
+# Single-stage NemotronHAudexForConditionalGeneration on
+# checkpoint_folder_full. Hybrid Mamba: no prefix caching. Pass this yaml
+# explicitly (see audex_tts_30b.yaml).
+pipeline: audex_thinker_only
+async_chunk: false
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.9
+    trust_remote_code: true
+    enable_prefix_caching: false
+    max_num_batched_tokens: 8192
+    # 30 clips x 750 sound embeddings (900 s cap) + text headroom.
+    max_model_len: 32768
+    dtype: bfloat16
+    devices: "0"
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 256
+      seed: 0

vllm_omni/deploy/audex_tta.yaml

@@ -0,0 +1,69 @@
+# Audex (Nemotron-Labs-Audex-2B) text-to-audio deploy: TTA thinker → XCodec1.
+#
+# Notable choices:
+#   * Sync full-payload (async_chunk: false): XCodec1 is a CNN codec the
+#     official flow decodes over the full sequence; there is no streaming
+#     reference and chunked decode risks boundary artifacts.
+#   * Stage 0 sampling defaults follow the official TTA script (temperature
+#     1.0, top_k 80; stop on <audiogen_end>; codec cap 4000 tokens = 1000
+#     frames generated, decode truncates at 500 frames ≈ 10 s).
+#   * Classifier-free guidance is effectively mandatory for TTA quality; the
+#     official default is cfg_scale 3.0. The RVQ phase mask is required for
+#     validity — both logits processors are registered here and driven per
+#     request via extra_args (cfg_* and tta_rvq).
+#   * Stage 1's model is the EXTERNAL hf-audio XCodec1 checkpoint (override
+#     with a local path here or via XCODEC1_PATH); its tokenizer is borrowed
+#     from the thinker checkpoint automatically.
+pipeline: audex_tta
+async_chunk: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      audiocodec_token_offset: 196613
+      audiocodec_vocab_size: 8192
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.4
+    trust_remote_code: true
+    enable_prefix_caching: true
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 8192
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+      - "vllm_omni.model_executor.models.audex.tta:TTARVQPhaseMaskLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 1.0
+      top_k: 80
+      max_tokens: 4200
+      seed: 0
+
+  - stage_id: 1
+    model: hf-audio/xcodec-hubert-general-balanced
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/deploy/audex_tta_30b.yaml

@@ -0,0 +1,57 @@
+# Audex 30B-A3B text-to-audio deploy: TTA thinker → XCodec1.
+# See audex_tta.yaml for the task notes and audex_tts_30b.yaml for the
+# 30B sizing notes (hybrid Mamba: no prefix caching; single-H100 layout,
+# TP2 fallback; pass this yaml explicitly).
+pipeline: audex_tta
+async_chunk: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      audiocodec_token_offset: 196613
+      audiocodec_vocab_size: 8192
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.85
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 8192
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+      - "vllm_omni.model_executor.models.audex.tta:TTARVQPhaseMaskLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 1.0
+      top_k: 80
+      max_tokens: 4200
+      seed: 0
+
+  - stage_id: 1
+    model: hf-audio/xcodec-hubert-general-balanced
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.08
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/deploy/audex_tts.yaml

@@ -0,0 +1,83 @@
+# Audex (Nemotron-Labs-Audex-2B) TTS deploy: thinker → code2wav via
+# shared-memory chunk streaming. Verified on 1x H100.
+#
+# Notable choices:
+#   * The code2wav decoder is a stateful streaming session per request — no
+#     left-context replay; chunks carry only new frames, so
+#     codec_left_context_frames stays 0.
+#   * codec_token_offset / codec_vocab_size mirror the checkpoint tokenizer's
+#     <speechcodec_N> block (pinned by a unit test).
+#   * Stage 0 sampling defaults follow the official TTS script
+#     (temperature 0.1, top_k 80, max_tokens 2048).
+#   * Classifier-free guidance: the stage-0 logits processor + scheduler
+#     pair patches are always installed but idle until a request carries
+#     extra_args.cfg_scale > 1.0 (serving: extra_params.cfg_scale; offline:
+#     --cfg-scale). Guided requests decode at temperature 0.05 (CFG
+#     sharpens the distribution; measured en-24 cfg1.5 CER: 6.87% at
+#     temp 0.05 vs 7.31% at the unguided 0.1 — baseline 7.24%).
+#     Every guided request adds an unconditional companion sequence in
+#     stage 0 — at max_num_seqs 8 that means at most 4 guided requests
+#     decode concurrently; raise max_num_seqs if CFG-heavy traffic needs
+#     more parallelism.
+pipeline: audex_tts
+async_chunk: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      codec_chunk_frames: 25
+      # Emit the first chunk early for a fast first audible sample.
+      initial_codec_chunk_frames: 1
+      codec_token_offset: 131077
+      codec_vocab_size: 65536
+      # Decoder-internal emit granularity (frames per decode call).
+      decoder_emit_chunk_frames: 5
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.4
+    trust_remote_code: true
+    enable_prefix_caching: true
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 4096
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.1
+      top_k: 80
+      max_tokens: 2048
+      seed: 0
+
+  - stage_id: 1
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    # Sync scheduling: with async scheduling the runner builds per-request
+    # payloads against the NEXT step's input batch, tripping the one-entry-
+    # per-request check when chunk arrival staggers requests.
+    async_scheduling: false
+    # async_chunk appends codec frames per step; max_model_len must cover the
+    # accumulated stream (max_tokens 2048 codes plus prompt headroom).
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/deploy/audex_tts_30b.yaml

@@ -0,0 +1,76 @@
+# Audex 30B-A3B (nvidia/Nemotron-Labs-Audex-30B-A3B) TTS deploy.
+#
+# Same pipeline/topology/token space as the 2B (audex_tts.yaml); the thinker
+# is the hybrid Mamba+MoE NemotronHForCausalLM (~60 GB bf16), so:
+#   * enable_prefix_caching: false — hybrid Mamba stage.
+#   * Single H100 80GB layout: thinker 0.85 + decoder 0.08 share GPU 0.
+#     If weights + state do not fit (long sequences OOM), set
+#     tensor_parallel_size: 2 on stage 0 and move stage 1 to devices "2".
+#   * CUDA graphs stay ON for stage 0: the hybrid Mamba+MoE thinker
+#     captures cleanly and review benchmarks measured ~7x lower TTS RTF
+#     vs enforce_eager (0.30 vs 2.08 single-stream), on par with the 2B
+#     and NVIDIA's reference demo.
+#   * Classifier-free guidance quality (en-24, self-transcribed ASR CER):
+#     baseline cfg=1.0 -> 0.72%; guided cfg=1.5 at temperature 0.1 -> 0.72%
+#     (gate: guided <= baseline, PASS). Note the 30B's best guided
+#     temperature is 0.1 (the yaml default) — unlike the 2B's 0.05, which
+#     measured 0.79% here. The serving layer currently injects 0.05 for
+#     guided requests on all sizes; acceptable (still within the anchor
+#     +2pp standing gate) but a size-aware refinement is a known follow-up.
+# ALWAYS pass this yaml explicitly (--stage-configs-path): bare serve of the
+# 30B repo auto-detects the shared model_type and lands on the 2B-tuned yaml.
+pipeline: audex_tts
+async_chunk: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      codec_chunk_frames: 25
+      initial_codec_chunk_frames: 1
+      codec_token_offset: 131077
+      codec_vocab_size: 65536
+      decoder_emit_chunk_frames: 5
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.85
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 4096
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.1
+      top_k: 80
+      max_tokens: 2048
+      seed: 0
+
+  - stage_id: 1
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.08
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/engine/async_omni_engine.py

@@ -30,7 +30,7 @@ from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.input_processor import InputProcessor
 
-from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override
 from vllm_omni.config.stage_config import (
     DuplexSessionRuntimeConfig,
     load_deploy_config,
@@ -151,7 +151,7 @@ class AsyncOmniEngine:
         transfer_emitter: Any = None,
         log_stats: bool = False,
         tokenizer: str | None = None,
-        trust_remote_code: bool = False,
+        trust_remote_code: bool | None = None,
         **kwargs: Any,
     ) -> None:
         self.model = model
@@ -230,7 +230,11 @@ class AsyncOmniEngine:
         if deploy_config_path is not None:
             self.duplex_session_config = load_deploy_config(deploy_config_path).duplex_session
 
-        kwargs["trust_remote_code"] = trust_remote_code
+        # Tri-state: None means "not specified" — the deploy yaml's per-stage
+        # trust_remote_code stays in effect. An explicit True/False here is a
+        # global override (precedence: caller > deploy yaml > default False);
+        # the merge rule lives in with_trust_remote_code_override.
+        kwargs = with_trust_remote_code_override(kwargs, trust_remote_code)
         self.config_path, self.stage_configs = self._resolve_stage_configs(
             model,
             kwargs,
@@ -786,6 +790,11 @@ class AsyncOmniEngine:
                 supported_tasks=self.supported_tasks,
             )
             request.external_req_id = cid
+            # Companions are stage-0-final: without this worker-visible tag
+            # the AR runner would still emit downstream connector payloads
+            # for them (the orchestrator-side final_stage_id=0 registration
+            # alone does not reach the worker).
+            request = apply_omni_final_stage_metadata(request, 0)
 
             # Registration of this companion on stage-0's output processor is
             # deferred to Orchestrator._handle_add_companion, which routes
@@ -1142,7 +1151,7 @@ class AsyncOmniEngine:
         model: str,
         kwargs: dict[str, Any],
         *,
-        trust_remote_code: bool,
+        trust_remote_code: bool | None,
     ) -> tuple[str, list[Any]]:
         """Resolve stage configs and inject defaults shared by orchestrator/headless."""
 

vllm_omni/engine/stage_engine_core_proc.py

@@ -30,7 +30,10 @@ from vllm.v1.engine.utils import (
 
 from vllm_omni.distributed.omni_coordinator import create_stage_coord_client
 from vllm_omni.engine import OmniEngineCoreRequest
-from vllm_omni.engine.stage_init_utils import set_death_signal
+from vllm_omni.engine.stage_init_utils import (
+    maybe_apply_audex_cfg_patches,
+    set_death_signal,
+)
 
 logger = init_logger(__name__)
 
@@ -131,6 +134,10 @@ class StageEngineCoreProc(EngineCoreProc):
                 _vllm_engine_core_module.EngineCoreRequest,
             )
 
+            # Audex CFG scheduler patches must land before EngineCore builds
+            # its Scheduler; gated on the stage's logits_processors config.
+            maybe_apply_audex_cfg_patches(kwargs.get("vllm_config"))
+
             engine_core = StageEngineCoreProc(
                 *args,
                 engine_index=dp_rank,

vllm_omni/engine/stage_init_utils.py

@@ -767,11 +767,53 @@ def build_engine_args_dict(
     stage_id = stage_config.stage_id
 
     engine_args_dict = _to_dict(engine_args)
+    pipeline_model_root = model
     model = engine_args_dict.pop("model", None) or model
     stage_defines_tokenizer = (
         engine_args_dict.get("tokenizer") is not None or engine_args_dict.get("tokenizer_subdir") is not None
     )
+    audex_stage = str(engine_args_dict.get("model_stage") or "")
+    if audex_stage == "audex_xcodec":
+        # TTA stage 1 decodes with the external XCodec1 checkpoint, not a
+        # subfolder of the Audex repo: the source comes from XCODEC1_PATH,
+        # the stage yaml's own ``model`` entry, or the default HF repo — the
+        # pipeline-level model (the Audex root) is never a valid source.
+        from vllm_omni.model_executor.models.audex.checkpoint import (
+            ensure_audex_snapshot,
+            ensure_xcodec1_snapshot,
+        )
+
+        stage_model = None if model == pipeline_model_root else model
+        model = ensure_xcodec1_snapshot(os.environ.get("XCODEC1_PATH") or stage_model)
+        if not stage_defines_tokenizer:
+            # XCodec1 ships no tokenizer files; borrow the thinker's (same
+            # workaround as the TTS decoder stage).
+            audex_root = ensure_audex_snapshot(pipeline_model_root, profile="tta")
+            engine_args_dict["tokenizer"] = os.path.join(audex_root, "checkpoint_folder_audiogen")
+            stage_defines_tokenizer = True
+    elif audex_stage.startswith("audex"):
+        # Audex users pass the HF repo ROOT; make sure the required snapshot
+        # subset exists locally BEFORE subdir resolution, otherwise the
+        # subdirs get joined onto the raw repo id on a fresh cache. The
+        # profile keeps TTS-only deployments from pulling the full-checkpoint
+        # extras.
+        from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
+
+        if audex_stage == "audex_omni":
+            audex_profile = "full"
+        elif audex_stage == "audex_tta_thinker":
+            audex_profile = "tta"
+        else:
+            audex_profile = "tts"
+        model = ensure_audex_snapshot(model, profile=audex_profile)
     model = _resolve_model_tokenizer_paths(model, engine_args_dict)
+    if engine_args_dict.get("model_stage") in ("audex_thinker", "audex_tta_thinker"):
+        # Audex ships its thinker weights deduplicated into a sibling folder;
+        # replicate the official prepare-script symlink on first use. The TTA
+        # thinker loads the same checkpoint_folder_audiogen checkpoint.
+        from vllm_omni.model_executor.models.audex.checkpoint import ensure_audiogen_weights
+
+        ensure_audiogen_weights(model)
     apply_cli_tokenizer(
         engine_args_dict,
         cli_tokenizer=cli_tokenizer,
@@ -1197,3 +1239,26 @@ def initialize_diffusion_stage(
 
     od_config = build_diffusion_config(model, stage_cfg, metadata)
     return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline)
+
+
+def maybe_apply_audex_cfg_patches(vllm_config: Any) -> None:
+    """Install the Audex CFG scheduler patches for CFG-configured engines.
+
+    Must run in the engine-core process BEFORE ``Scheduler`` is constructed:
+    the patch wraps ``Scheduler.__init__`` to add the pair registry, so a
+    scheduler built earlier would never become pair-aware. Gated on the
+    engine's ``logits_processors`` so non-CFG stages stay untouched.
+    """
+    model_config = getattr(vllm_config, "model_config", None)
+    processors = getattr(model_config, "logits_processors", None) or []
+    if not any("AudexCFGLogitsProcessor" in getattr(proc, "__name__", str(proc)) for proc in processors):
+        return
+
+    from vllm_omni.model_executor.models.audex.cfg import apply_cfg_patches
+
+    apply_cfg_patches()
+
+    from vllm.v1.core.sched.scheduler import Scheduler
+
+    if not getattr(Scheduler.schedule, "_audex_cfg_patched", False):
+        raise RuntimeError("Audex CFG scheduler patches failed to install before Scheduler construction")

vllm_omni/entrypoints/cli/serve.py

@@ -852,7 +852,9 @@ def run_headless(args: TrackingNamespace) -> None:
         model,
         stage_configs_path,
         args_dict,
-        trust_remote_code=bool(getattr(args, "trust_remote_code", False)),
+        # store_true cannot express an explicit False: absent maps to None
+        # ("not specified") so the deploy yaml's per-stage value applies.
+        trust_remote_code=getattr(args, "trust_remote_code", None) or None,
         deploy_config_path=args_dict.get("deploy_config"),
         stage_overrides=stage_overrides,
         strategy_config_path=args_dict.get("strategy_config"),

vllm_omni/entrypoints/openai/serving_speech.py

@@ -100,6 +100,10 @@ _HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"}
 _GLM_TTS_MODEL_STAGES = {"glm_tts"}
 _STEP_AUDIO2_TTS_MODEL_STAGES = {"step_audio2_thinker"}
 _INDEXTTS2_TTS_MODEL_STAGES = {"indextts2_talker"}
+# audex_omni covers the audex_s2s S2S deployment, whose
+# TTS pass uses the same /v1/audio/speech surface.
+_AUDEX_TTS_MODEL_STAGES = {"audex_thinker", "audex_omni"}
+_AUDEX_TTA_MODEL_STAGES = {"audex_tta_thinker"}
 _TTS_MODEL_STAGES: set[str] = (
     _VOXTRAL_TTS_MODEL_STAGES
     | _QWEN3_TTS_MODEL_STAGES
@@ -117,6 +121,8 @@ _TTS_MODEL_STAGES: set[str] = (
     | _GLM_TTS_MODEL_STAGES
     | _STEP_AUDIO2_TTS_MODEL_STAGES
     | _INDEXTTS2_TTS_MODEL_STAGES
+    | _AUDEX_TTS_MODEL_STAGES
+    | _AUDEX_TTA_MODEL_STAGES
 )
 _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
     "fish_tts",
@@ -127,7 +133,13 @@ _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
     "higgs_audio_v2",
     "higgs_audio_v3",
     "indextts2",
+    "audex",
+    "audex_tta",
 }
+# Audex contract: zero-codec / invalid generations arrive as empty terminal
+# payloads and must fail the request, never serialize as a successful empty
+# WAV. Covers both the TTS ("audex") and TTA ("audex_tta") pipelines.
+_AUDEX_NO_AUDIO_GUARD_MODEL_TYPES = frozenset({"audex", "audex_tta"})
 _TTS_LANGUAGES = frozenset(
     {
         "Auto",
@@ -493,6 +505,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         self._tts_tokenizer = None
         self._voxcpm2_tokenizer = None
+        self._audex_tokenizer = None
+        self._audex_tta_rvq = None
         self._voxcpm2_split_map: dict[int, list[int]] = {}
 
         logger.info("Loaded %d supported speakers: %s", len(self.supported_speakers), sorted(self.supported_speakers))
@@ -673,12 +687,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
     def _find_tts_stage(self):
         """Find and return the TTS stage config, or None if not found."""
+        all_stages = {getattr(stage.engine_args, "model_stage", None) for stage in self.engine_client.stage_configs}
         for stage in self.engine_client.stage_configs:
             engine_args = stage.engine_args
             model_stage = engine_args.model_stage
             model_arch = getattr(engine_args, "model_arch", None)
             worker_type = getattr(engine_args, "worker_type", None)
             if model_stage in _TTS_MODEL_STAGES:
+                # The audio-capable Audex thinker is only speech-capable when
+                # deployed WITH the speech decoder (audex_s2s);
+                # the thinker-only deployment is text-final and must not
+                # accept /v1/audio/speech requests.
+                if model_stage == "audex_omni" and "audex_code2wav" not in all_stages:
+                    continue
                 return stage
             # Ming dense identifies its AR entry stage by architecture because
             # it does not use a dedicated TTS model_stage value.
@@ -729,6 +750,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return "step_audio2"
         if model_stage in _INDEXTTS2_TTS_MODEL_STAGES:
             return "indextts2"
+        if model_stage in _AUDEX_TTS_MODEL_STAGES:
+            return "audex"
+        if model_stage in _AUDEX_TTA_MODEL_STAGES:
+            return "audex_tta"
         return None
 
     def _get_custom_voice_dir(self) -> str | None:
@@ -2318,6 +2343,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     )
                     if chunk_np.ndim > 1:
                         chunk_np = chunk_np.squeeze()
+                    if self._tts_model_type in _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES and int(np.size(chunk_np)) == 0:
+                        # Zero-size chunks must not emit a WAV header or count
+                        # as first audio; the post-loop guard below needs to
+                        # see an audio-less stream to fail the request.
+                        continue
                     # For WAV format, emit header before first audio chunk
                     if response_format == "wav" and first_chunk:
                         # Assert that sample rate has been set from chunk metadata (not just default)
@@ -2349,6 +2379,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         yield audio_bytes, sample_rate_val
                     else:
                         yield audio_bytes
+            if self._tts_model_type in _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES and first_audio_chunk_s is None:
+                # Audex contract: zero codec tokens must abort the stream, not
+                # complete it cleanly with zero audio bytes.
+                raise ValueError("Audex produced no audio (the thinker emitted zero or invalid codec tokens)")
             self._mark_ref_audio_artifact_ready_for_request(request_id)
             artifact_ready = True
             total_ms = (time.perf_counter() - stream_start_s) * 1000.0
@@ -2822,6 +2856,140 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         )
         return {"prompt_token_ids": prompt_ids}
 
+    def _get_audex_tokenizer(self):
+        """Load (once per process) the Audex thinker tokenizer for CFG/TTA prompts.
+
+        The folder follows the serving stage: the full-pipeline thinker
+        (``audex_omni``) tokenizes with ``checkpoint_folder_full``; the TTS
+        and TTA thinkers use ``checkpoint_folder_audiogen``.
+        """
+        if self._audex_tokenizer is None:
+            from transformers import AutoTokenizer
+
+            from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
+
+            stage = getattr(getattr(self._tts_stage, "engine_args", None), "model_stage", None)
+            if stage == "audex_omni":
+                profile, folder = "full", "checkpoint_folder_full"
+            elif stage == "audex_tta_thinker":
+                profile, folder = "tta", "checkpoint_folder_audiogen"
+            else:
+                profile, folder = "tts", "checkpoint_folder_audiogen"
+            # The serving model_config may already hold the RESOLVED per-stage
+            # checkpoint folder (stage init joins model_subdir onto the
+            # snapshot root); joining again would produce a non-existent path
+            # that transformers then mistakes for a repo id.
+            model_path = os.path.normpath(self.engine_client.model_config.model)
+            if os.path.basename(model_path).startswith("checkpoint_folder"):
+                root = os.path.dirname(model_path)
+            else:
+                root = ensure_audex_snapshot(model_path, profile=profile)
+            self._audex_tokenizer = AutoTokenizer.from_pretrained(os.path.join(root, folder))
+        return self._audex_tokenizer
+
+    def _inject_audex_tta_args(self, request_id: str, prompt: Any, stage0_params: Any) -> Any:
+        """Complete the Audex TTA contract on a CLONE of the stage-0 params.
+
+        Always attaches the RVQ phase-mask contract (required for token
+        validity) and the CFG pair contract at the official default scale
+        3.0 unless the caller passed ``cfg_scale`` (1.0 disables guidance).
+
+        Returns the cloned params; the input is never mutated. On the
+        no-extra_params online path the input is an element of the engine's
+        SHARED ``default_sampling_params_list`` — writing per-request pair
+        state (``cfg_pair_id``/``cfg_null_prompt``) into it would race
+        concurrent requests and leak stale CFG metadata into later ones.
+        """
+        import copy
+
+        from vllm_omni.model_executor.models.audex.prompt import build_tta_null_prompt
+        from vllm_omni.model_executor.models.audex.tta import build_tta_phase_token_ids
+
+        cond_prompt = prompt.get("prompt") if isinstance(prompt, dict) else prompt
+        if not isinstance(cond_prompt, str) or not cond_prompt:
+            raise ValueError("Audex TTA requires the adapter-built caption prompt")
+
+        stage0_params = copy.deepcopy(stage0_params)
+        tokenizer = self._get_audex_tokenizer()
+        if self._audex_tta_rvq is None:
+            phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer)
+            self._audex_tta_rvq = {
+                "phase_token_ids": phase_token_ids,
+                "start_tid": start_tid,
+                "end_tid": end_tid,
+                # Official generation cap (decode truncates at 500 frames).
+                "codec_cap": 4000,
+                # The TTA prompt ends with <audiogen_start>.
+                "start_in_prompt": True,
+            }
+
+        extra_args = getattr(stage0_params, "extra_args", None)
+        if extra_args is None:
+            extra_args = {}
+            stage0_params.extra_args = extra_args
+        extra_args["tta_rvq"] = self._audex_tta_rvq
+
+        cfg_scale = extra_args.get("cfg_scale")
+        cfg_scale = 3.0 if cfg_scale is None else float(cfg_scale)
+        if cfg_scale <= 1.0:
+            extra_args.pop("cfg_scale", None)
+            return stage0_params
+        extra_args.update(
+            {
+                "cfg_scale": cfg_scale,
+                "cfg_role": "cond",
+                "cfg_pair_id": request_id,
+                "cfg_null_prompt": build_tta_null_prompt(cond_prompt, tokenizer),
+            }
+        )
+        return stage0_params
+
+    def _inject_audex_cfg_pair_args(self, request_id: str, prompt: Any, stage0_params: Any) -> Any:
+        """Complete the Audex CFG contract on a CLONE of the stage-0 params.
+
+        The adapter validated ``cfg_scale``; here (where the final request id
+        exists) the cond role, pair id, and the length-matched null prompt are
+        attached. ``cfg_scale`` absent/1.0 returns a clone value-identical to
+        the non-CFG flow (the key is dropped rather than forwarded).
+
+        Returns the cloned params; the input is never mutated (it may be an
+        element of the engine's shared ``default_sampling_params_list`` —
+        see ``_inject_audex_tta_args``).
+        """
+        import copy
+
+        stage0_params = copy.deepcopy(stage0_params)
+        extra_args = getattr(stage0_params, "extra_args", None)
+        if not extra_args or extra_args.get("cfg_scale") is None:
+            if extra_args:
+                extra_args.pop("cfg_scale", None)
+            return stage0_params
+        cfg_scale = float(extra_args["cfg_scale"])
+        if cfg_scale <= 1.0:
+            extra_args.pop("cfg_scale", None)
+            return stage0_params
+
+        from vllm_omni.model_executor.models.audex.prompt import build_null_prompt
+
+        cond_prompt = prompt.get("prompt") if isinstance(prompt, dict) else prompt
+        if not isinstance(cond_prompt, str) or not cond_prompt:
+            raise ValueError("Audex CFG requires the adapter-built text prompt")
+        null_prompt = build_null_prompt(cond_prompt, self._get_audex_tokenizer())
+        extra_args.update(
+            {
+                "cfg_scale": cfg_scale,
+                "cfg_role": "cond",
+                "cfg_pair_id": request_id,
+                "cfg_null_prompt": null_prompt,
+            }
+        )
+        # Guided decoding sharpens the distribution; the unguided default
+        # temperature (0.1) adds excess sampling noise under CFG. Measured
+        # on the en-24 gate corpus at cfg 1.5: temp 0.1 -> CER 7.31%,
+        # temp 0.05 -> CER 6.87% (unguided baseline 7.24%).
+        stage0_params.temperature = 0.05
+        return stage0_params
+
     def _apply_cosyvoice3_dynamic_tokens(
         self,
         sampling_params_list: list,
@@ -3205,6 +3373,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             sampling_params_list[0].extra_args.update(request.extra_params)
             logger.info("Applied extra_params: %s", request.extra_params)
 
+        # Audex CFG: turn an adapter-validated cfg_scale into the engine-side
+        # pair contract. This must run here (not in the adapter) because the
+        # pair id is the final request_id, which the adapter never sees.
+        # The injectors return a CLONE: without extra_params above,
+        # sampling_params_list still holds the engine's shared default params,
+        # which must never carry per-request CFG state.
+        if self._tts_model_type == "audex" and sampling_params_list:
+            sampling_params_list[0] = self._inject_audex_cfg_pair_args(request_id, prompt, sampling_params_list[0])
+        elif self._tts_model_type == "audex_tta" and sampling_params_list:
+            sampling_params_list[0] = self._inject_audex_tta_args(request_id, prompt, sampling_params_list[0])
+
         # Some TTS model defaults come from deploy YAML. Their AR
         # generation length is controlled by SamplingParams.max_tokens, so only
         # override it when the caller explicitly requests max_new_tokens.
@@ -3405,6 +3584,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             if audio_tensor.ndim > 1:
                 audio_tensor = audio_tensor.squeeze()
 
+            if self._tts_model_type in _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES and int(np.size(audio_tensor)) == 0:
+                # Audex contract: zero codec tokens must fail the request, not
+                # serialize as an empty-but-successful WAV.
+                raise ValueError("Audex produced no audio (the thinker emitted zero or invalid codec tokens)")
+
             audio_obj = CreateAudio(
                 audio_tensor=audio_tensor,
                 sample_rate=sample_rate,

vllm_omni/entrypoints/openai/tts_adapters/__init__.py

@@ -51,6 +51,8 @@ def resolve_adapter(model_type: str | None) -> type[TTSModelAdapter] | None:
 # Import adapter modules for their registration side effects. Keep at the bottom
 # so the registry helpers above are defined first.
 from vllm_omni.entrypoints.openai.tts_adapters import (  # noqa: E402,F401
+    audex,
+    audex_tta,
     cosyvoice3,
     covo_audio,
     fish_speech,

vllm_omni/entrypoints/openai/tts_adapters/audex.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Audex (Nemotron-Labs-Audex-2B) TTS serving adapter."""
+
+from typing import TYPE_CHECKING
+
+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_executor.models.audex.prompt import build_cond_prompt
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+# Classifier-free guidance strength accepted on /v1/audio/speech via
+# extra_params. 1.0 disables guidance (identical to omitting it); the
+# official TTS quality setting is 1.5.
+AUDEX_CFG_SCALE_MIN = 1.0
+AUDEX_CFG_SCALE_MAX = 10.0
+
+# Internal CFG pair-plumbing keys; injected by the serving layer once the
+# request id exists, never accepted from callers.
+_AUDEX_INTERNAL_CFG_KEYS = ("cfg_role", "cfg_pair_id", "cfg_null_prompt")
+
+
+@register_tts_adapter
+class AudexAdapter(ARTTSAdapter):
+    """Plain English TTS: single built-in voice, no reference audio, optional CFG."""
+
+    stage_keys = frozenset({"audex_thinker", "audex_omni"})
+    name = "audex"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        if not request.input or not request.input.strip():
+            return "Audex TTS requires non-empty input text"
+        voice = (request.voice or "").strip().lower()
+        if voice not in ("", "default"):
+            return (
+                f"Audex has a single built-in voice and no voice cloning; got voice={request.voice!r}. "
+                "Omit 'voice' or pass 'default'."
+            )
+        if request.ref_audio is not None:
+            return "Audex does not support reference audio (no voice cloning)."
+        extra = request.extra_params or {}
+        for key in _AUDEX_INTERNAL_CFG_KEYS:
+            if key in extra:
+                return f"extra_params.{key} is managed internally by the server; only cfg_scale may be set."
+        cfg_scale = extra.get("cfg_scale")
+        if cfg_scale is not None:
+            try:
+                cfg_value = float(cfg_scale)
+            except (TypeError, ValueError):
+                return (
+                    f"extra_params.cfg_scale must be a number in "
+                    f"[{AUDEX_CFG_SCALE_MIN}, {AUDEX_CFG_SCALE_MAX}]; got {cfg_scale!r}."
+                )
+            if not (AUDEX_CFG_SCALE_MIN <= cfg_value <= AUDEX_CFG_SCALE_MAX):
+                return (
+                    f"extra_params.cfg_scale must be within [{AUDEX_CFG_SCALE_MIN}, {AUDEX_CFG_SCALE_MAX}]; "
+                    f"got {cfg_scale!r}. 1.0 disables guidance; 1.5 is the recommended quality setting."
+                )
+        return None
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        prompt = {"prompt": build_cond_prompt(request.input)}
+        return PreparedRequest(prompt=prompt, tts_params={}, model_type="audex")

vllm_omni/entrypoints/openai/tts_adapters/audex_tta.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Audex (Nemotron-Labs-Audex-2B) text-to-audio serving adapter.
+
+Caption → general audio through the ``audex_tta`` pipeline.
+The adapter builds the TTA prompt; the serving layer completes the request
+after the final request id exists: the RVQ phase-mask contract
+(``extra_args["tta_rvq"]``) and the CFG pair contract (official default
+scale 3.0 — guidance is effectively mandatory for TTA quality).
+"""
+
+from typing import TYPE_CHECKING
+
+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_executor.models.audex.prompt import build_tta_cond_prompt
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+AUDEX_TTA_DEFAULT_CFG_SCALE = 3.0
+AUDEX_TTA_CFG_SCALE_MIN = 1.0
+AUDEX_TTA_CFG_SCALE_MAX = 10.0
+
+# Internal plumbing keys; injected by the serving layer, never accepted
+# from callers.
+_INTERNAL_KEYS = ("cfg_role", "cfg_pair_id", "cfg_null_prompt", "tta_rvq")
+
+
+@register_tts_adapter
+class AudexTTAAdapter(ARTTSAdapter):
+    """Caption-conditioned general audio: no voices, no reference audio."""
+
+    stage_keys = frozenset({"audex_tta_thinker"})
+    name = "audex_tta"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        if not request.input or not request.input.strip():
+            return "Audex TTA requires a non-empty caption"
+        voice = (request.voice or "").strip().lower()
+        if voice not in ("", "default"):
+            return f"Audex TTA generates general audio and has no voices; got voice={request.voice!r}."
+        if request.ref_audio is not None:
+            return "Audex TTA does not support reference audio."
+        extra = request.extra_params or {}
+        for key in _INTERNAL_KEYS:
+            if key in extra:
+                return f"extra_params.{key} is managed internally by the server; only cfg_scale may be set."
+        cfg_scale = extra.get("cfg_scale")
+        if cfg_scale is not None:
+            try:
+                cfg_value = float(cfg_scale)
+            except (TypeError, ValueError):
+                return (
+                    f"extra_params.cfg_scale must be a number in "
+                    f"[{AUDEX_TTA_CFG_SCALE_MIN}, {AUDEX_TTA_CFG_SCALE_MAX}]; got {cfg_scale!r}."
+                )
+            if not (AUDEX_TTA_CFG_SCALE_MIN <= cfg_value <= AUDEX_TTA_CFG_SCALE_MAX):
+                return (
+                    f"extra_params.cfg_scale must be within "
+                    f"[{AUDEX_TTA_CFG_SCALE_MIN}, {AUDEX_TTA_CFG_SCALE_MAX}]; got {cfg_scale!r}. "
+                    f"The official TTA setting is {AUDEX_TTA_DEFAULT_CFG_SCALE}; 1.0 disables guidance."
+                )
+        return None
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        prompt = {"prompt": build_tta_cond_prompt(request.input)}
+        return PreparedRequest(prompt=prompt, tts_params={}, model_type="audex_tta")

vllm_omni/entrypoints/utils.py

@@ -12,7 +12,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.transformers_utils.config import get_config, get_hf_file_to_dict
 from vllm.transformers_utils.repo_utils import file_or_path_exists
 
-from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override
 from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs
 from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
 from vllm_omni.entrypoints.stage_utils import _to_dict
@@ -103,12 +103,7 @@ def _filter_dict_like_object(obj: dict | Any) -> dict:
             return True
         return isinstance(
             value,
-            (
-                types.FunctionType,
-                types.MethodType,
-                types.BuiltinFunctionType,
-                types.BuiltinMethodType,
-            ),
+            types.FunctionType | types.MethodType | types.BuiltinFunctionType | types.BuiltinMethodType,
         )
 
     result = {}
@@ -193,10 +188,10 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
     if callable(obj):
         return None
     # Handle lists and tuples (recurse into items)
-    if isinstance(obj, (list, tuple)):
+    if isinstance(obj, list | tuple):
         return type(obj)(_convert_dataclasses_to_dict(item) for item in obj if not callable(item))
     # Try to convert any dict-like object (has keys/values methods) to dict
-    if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, (str, bytes)):
+    if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, str | bytes):
         try:
             return _filter_dict_like_object(obj)
         except (TypeError, ValueError, AttributeError):
@@ -315,7 +310,7 @@ def resolve_model_config_path(model: str) -> str:
 def load_stage_configs_from_model(
     model: str,
     *,
-    trust_remote_code: bool,
+    trust_remote_code: bool | None,
     base_engine_args: dict | None = None,
     deploy_config_path: str | None = None,
     stage_overrides: dict[str, dict[str, Any]] | None = None,
@@ -349,7 +344,12 @@ def load_stage_configs_from_model(
         base_engine_args = {}
 
     cli_overrides = _convert_dataclasses_to_dict(dict(base_engine_args))
-    cli_overrides["trust_remote_code"] = trust_remote_code
+    # A False inherited from the engine-args dump is the store_true flag's
+    # default, not an explicit choice — drop it so only the tri-state
+    # parameter below decides (see with_trust_remote_code_override).
+    if not cli_overrides.get("trust_remote_code"):
+        cli_overrides.pop("trust_remote_code", None)
+    cli_overrides = with_trust_remote_code_override(cli_overrides, trust_remote_code)
     if stage_overrides:
         for stage_id_str, overrides in stage_overrides.items():
             for key, val in overrides.items():
@@ -559,7 +559,7 @@ def load_and_resolve_stage_configs(
     stage_configs_path: str | None,
     kwargs: dict | None,
     *,
-    trust_remote_code: bool,
+    trust_remote_code: bool | None,
     default_stage_cfg_factory: Any = None,
     deploy_config_path: str | None = None,
     stage_overrides: dict[str, dict[str, Any]] | None = None,
@@ -728,7 +728,7 @@ def filter_dataclass_kwargs(cls: Any, kwargs: dict) -> dict:
         if origin in (list, tuple, set):
             args = get_args(annotation)
             inner = args[0] if args else None
-            if isinstance(value, (list, tuple, set)):
+            if isinstance(value, list | tuple | set):
                 return type(value)(_filter_value(v, inner) for v in value)
             return value
 

vllm_omni/model_executor/models/audex/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/audex/audex_code2wav.py

@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Audex code2wav: streaming speech-codec-token → waveform stage.
+
+Wraps the vendored Audex causal speech decoder (see ``speech_decoder/``).
+Unlike the flow/vocoder stages of other TTS models, the decoder is natively
+stateful and streaming: one ``AudexCausalSpeechDecoderSession`` per request
+holds the decode cache, each engine step pushes only the newly received codec
+frames, and the session's 4-frame lookahead tail is drained by ``flush()``
+when the request finishes. The stage therefore never re-decodes left context
+and always returns delta audio.
+"""
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+from torch import nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.audex.speech_decoder import (
+    AudexCausalSpeechDecoderConfig,
+    AudexCausalSpeechDecoderModel,
+    AudexCausalSpeechDecoderSession,
+)
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+logger = init_logger(__name__)
+
+# Frames per decoder emit. 5 frames = 100 ms at 50 fps — the granularity the
+# official streaming demo uses; small enough for a fast first audible chunk.
+_DEFAULT_EMIT_CHUNK_FRAMES = 5
+
+
+def _meta_int(value: Any, default: int = 0) -> int:
+    if isinstance(value, torch.Tensor):
+        return int(value.reshape(-1)[0].item()) if value.numel() else default
+    if isinstance(value, list | tuple):
+        return _meta_int(value[0], default) if value else default
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return default
+
+
+def _meta_bool(value: Any) -> bool:
+    if isinstance(value, torch.Tensor):
+        return bool(value.reshape(-1)[0].item()) if value.numel() else False
+    if isinstance(value, list | tuple):
+        return _meta_bool(value[0]) if value else False
+    return bool(value)
+
+
+def _meta_str(value: Any) -> str | None:
+    if isinstance(value, list | tuple):
+        return _meta_str(value[0]) if value else None
+    if isinstance(value, str):
+        return value
+    return None
+
+
+def _codes_from_runtime_info(runtime_info: dict[str, Any] | None) -> list[int]:
+    """Extract this step's new codec frame ids from the connector payload."""
+    if not isinstance(runtime_info, dict):
+        return []
+    codes = runtime_info.get("codes")
+    if not isinstance(codes, dict):
+        return []
+    audio = codes.get("audio")
+    if isinstance(audio, torch.Tensor):
+        return audio.reshape(-1).tolist()
+    if isinstance(audio, list | tuple):
+        return [int(c) for c in audio]
+    return []
+
+
+class AudexCode2Wav(nn.Module):
+    """Stage-1 model for Audex TTS (GenerationModelRunner)."""
+
+    input_modalities = "audio"
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.model_path = vllm_config.model_config.model
+
+        self.have_multimodal_outputs = True
+        self.has_preprocess = False
+        self.has_postprocess = False
+        self.enable_update_additional_information = True
+        self.requires_raw_input_tokens = True
+
+        # Build the decoder skeleton from config so vLLM's memory profiler
+        # sees it at startup; weights load later in load_weights().
+        config = AudexCausalSpeechDecoderConfig.from_pretrained(self.model_path)
+        self.decoder = AudexCausalSpeechDecoderModel(config)
+        self.decoder.eval()
+        self._sample_rate = int(config.sample_rate)
+
+        extra_cfg = self._connector_extra_config()
+        self._emit_chunk_frames = _meta_int(extra_cfg.get("decoder_emit_chunk_frames"), _DEFAULT_EMIT_CHUNK_FRAMES)
+        if self._emit_chunk_frames <= 0:
+            raise ValueError(f"decoder_emit_chunk_frames must be positive, got {self._emit_chunk_frames}")
+
+        # One streaming session per in-flight request, keyed by the request id
+        # carried in each chunk payload's meta (equal to the stage-1 engine
+        # request id, so on_requests_finished can free aborted requests too).
+        self._sessions: dict[str, AudexCausalSpeechDecoderSession] = {}
+
+    def _connector_extra_config(self) -> dict[str, Any]:
+        connector_cfg = getattr(self.vllm_config.model_config, "stage_connector_config", None)
+        if isinstance(connector_cfg, dict):
+            extra = connector_cfg.get("extra", connector_cfg)
+            return extra if isinstance(extra, dict) else {}
+        extra = getattr(connector_cfg, "extra", None)
+        return extra if isinstance(extra, dict) else {}
+
+    # -------------------- vLLM model interface --------------------
+
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        # This stage ignores token embeddings; keep a stable dummy for the runner.
+        if input_ids.numel() == 0:
+            return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
+        return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
+
+    def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None:
+        return None
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        state_dict = dict(weights)
+        self.decoder.load_state_dict(state_dict, strict=True)
+        device = self.vllm_config.device_config.device
+        self.decoder.to(device=device, dtype=torch.float32)
+        return {f"decoder.{name}" for name in state_dict}
+
+    def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None:
+        """Free decoder sessions for finished/aborted requests."""
+        for req_id in finished_req_ids:
+            self._sessions.pop(req_id, None)
+
+    def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput:
+        if isinstance(model_outputs, OmniOutput):
+            return model_outputs
+        # The runner may unpack the NamedTuple in transit; rebuild it.
+        if isinstance(model_outputs, tuple) and len(model_outputs) == len(OmniOutput._fields):
+            return OmniOutput(*model_outputs)
+        raise TypeError(f"AudexCode2Wav expected OmniOutput, got {type(model_outputs)}")
+
+    # -------------------- decoding --------------------
+
+    def _session_for(self, req_id: str) -> AudexCausalSpeechDecoderSession:
+        session = self._sessions.get(req_id)
+        if session is None:
+            session = self.decoder.create_session(
+                chunk_frames=self._emit_chunk_frames,
+                return_numpy=False,
+            )
+            self._sessions[req_id] = session
+        return session
+
+    def _decode_request(self, req_id: str, codes: list[int], finished: bool) -> torch.Tensor:
+        """Push new frames into the request's session; return this step's delta audio."""
+        if finished and not codes and req_id not in self._sessions:
+            # Zero codec tokens for the whole request. Do NOT raise: an
+            # exception in forward() kills the engine core for every request.
+            # Returning empty audio lets the serving layer fail this request.
+            logger.error("Audex code2wav: request %s finished with no codec tokens; no audio generated", req_id)
+            return torch.empty(0, dtype=torch.float32)
+        chunks: list[torch.Tensor] = []
+        session = self._session_for(req_id)
+        with torch.inference_mode():
+            if codes:
+                for _sr, chunk in session.push([[int(c)] for c in codes]):
+                    chunks.append(chunk.reshape(-1))
+            if finished:
+                for _sr, chunk in session.flush():
+                    chunks.append(chunk.reshape(-1))
+        if finished:
+            self._sessions.pop(req_id, None)
+        if not chunks:
+            return torch.empty(0, dtype=torch.float32)
+        return torch.cat(chunks).to(torch.float32)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: Any = None,
+        inputs_embeds: torch.Tensor | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        runtime_infos = runtime_additional_information or []
+        # One output entry per scheduled request, even on steps where a
+        # request's chunk payload has not arrived yet.
+        seq_token_counts = kwargs.get("seq_token_counts")
+        if seq_token_counts is not None:
+            num_reqs = len(seq_token_counts)
+        elif input_ids is not None and input_ids.numel() > 0 and not runtime_infos:
+            num_reqs = 1
+        else:
+            num_reqs = len(runtime_infos)
+
+        audios: list[torch.Tensor] = []
+        sr_tensor = torch.tensor(self._sample_rate, dtype=torch.int32)
+        for i in range(num_reqs):
+            info = runtime_infos[i] if i < len(runtime_infos) else {}
+            info = info if isinstance(info, dict) else {}
+            meta = info.get("meta", {})
+            meta = meta if isinstance(meta, dict) else {}
+            req_id = _meta_str(meta.get("req_id"))
+            # async_chunk delivers ``stream_finished`` (the adapter strips
+            # ``finished``); the sync full-payload path delivers ``finished``.
+            finished = _meta_bool(meta.get("stream_finished")) or _meta_bool(meta.get("finished"))
+            codes = _codes_from_runtime_info(info)
+            if req_id is None:
+                if codes:
+                    raise ValueError("Audex code2wav chunk payload is missing meta.req_id")
+                audios.append(torch.empty(0, dtype=torch.float32))
+                continue
+            audios.append(self._decode_request(req_id, codes, finished))
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"model_outputs": audios, "sr": [sr_tensor] * num_reqs},
+        )

vllm_omni/model_executor/models/audex/audex_omni.py

@@ -0,0 +1,322 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0):
+#   inference_scripts_vllm/audioqa_scripts/audex_2b_vllm/modeling_audex_vllm.py
+#   inference_scripts_vllm/audioqa_scripts/audex_2b_vllm/audio_encoder.py
+"""Audex audio-understanding model (``checkpoint_folder_full``): S2T / audio QA.
+
+Mirrors the Qwen2-Audio pattern with the 2B dense backbone:
+  * ``language_model``  = the vllm-omni ``NemotronDenseForCausalLM`` port
+    (wrapped via ``init_vllm_registered_model``; resolved through the omni
+    model registry).
+  * ``audio_tower``     = transformers ``Qwen2AudioEncoder`` ("NV-Whisper");
+    a parameter-free avg pool maps 1500 → 750 frames per padded 30 s clip.
+  * ``audio_projector`` = RMSNorm → fc1 → relu² → fc2 (no bias), matching the
+    release ``NemotronHAudexProjector``.
+
+``embed_multimodal(...)`` runs encoder+projector and returns one embedding
+tensor per audio item; vLLM's ``SupportsMultiModal.embed_input_ids`` merges
+them into the ``<so_embedding>`` placeholder positions. The encoder runs in
+microbatches over clips so peak memory is bounded regardless of audio length
+(up to the 900 s / 30-clip cap enforced by the processor).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Annotated, Literal
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.config import VllmConfig
+from vllm.model_executor.models.interfaces import (
+    HasInnerState,
+    IsHybrid,
+    MultiModalEmbeddings,
+    SupportsMambaPrefixCaching,
+    SupportsMultiModal,
+    SupportsPP,
+)
+from vllm.model_executor.models.utils import (
+    init_vllm_registered_model,
+    maybe_prefix,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
+
+from vllm_omni.model_executor.models.audex.audex_omni_processing import (
+    AUDIO_ENCODER_MICRO_BATCH,
+    SOUND_TOKEN,
+    AudexDummyInputsBuilder,
+    AudexMultiModalProcessor,
+    AudexProcessingInfo,
+)
+
+
+class AudexProjector(nn.Module):
+    """Megatron sound_projection equivalent (matches HF NemotronHAudexProjector)."""
+
+    def __init__(
+        self,
+        audio_encoder_hidden_size: int,
+        intermediate_size: int,
+        text_hidden_size: int,
+        activation: str = "relu2",
+        norm_eps: float = 1e-5,
+    ) -> None:
+        super().__init__()
+        self.norm_weight = nn.Parameter(torch.ones(audio_encoder_hidden_size))
+        self.norm_eps = norm_eps
+        self.fc1 = nn.Linear(audio_encoder_hidden_size, intermediate_size, bias=False)
+        self.fc2 = nn.Linear(intermediate_size, text_hidden_size, bias=False)
+        self.activation = activation
+
+    def _rms_norm(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.norm_eps)
+        return (self.norm_weight.float() * x).to(dtype)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self._rms_norm(hidden_states)
+        hidden_states = self.fc1(hidden_states)
+        if self.activation == "relu2":
+            hidden_states = F.relu(hidden_states).pow(2)
+        elif self.activation == "gelu":
+            hidden_states = F.gelu(hidden_states)
+        else:
+            raise ValueError(f"Unsupported audio projector activation: {self.activation}")
+        return self.fc2(hidden_states)
+
+
+def _build_qwen2_audio_encoder(audio_config: dict) -> nn.Module:
+    from transformers.models.qwen2_audio.configuration_qwen2_audio import (
+        Qwen2AudioEncoderConfig,
+    )
+    from transformers.models.qwen2_audio.modeling_qwen2_audio import Qwen2AudioEncoder
+
+    cfg = Qwen2AudioEncoderConfig(**dict(audio_config))
+    return Qwen2AudioEncoder(cfg)
+
+
+class AudexAudioFeatureInputs(TensorSchema):
+    """
+    Dimensions:
+        - c: total number of 30s clips (flattened across audio items)
+        - b: number of audio items
+    """
+
+    type: Literal["audio_features"] = "audio_features"
+    input_features: Annotated[torch.Tensor, TensorShape("c", 128, 3000)]
+    audio_num_clips: Annotated[list[int], None] = None
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    AudexMultiModalProcessor,
+    info=AudexProcessingInfo,
+    dummy_inputs=AudexDummyInputsBuilder,
+)
+class NemotronDenseAudexForConditionalGeneration(
+    nn.Module,
+    SupportsMultiModal,
+    SupportsPP,
+):
+    _LM_ARCHITECTURE = "NemotronDenseForCausalLM"
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        if modality.startswith("audio"):
+            return SOUND_TOKEN
+        raise ValueError(f"Only audio modality is supported, got {modality}")
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        self.config = config
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                hf_config=config,
+                prefix=maybe_prefix(prefix, "language_model"),
+                architectures=[self._LM_ARCHITECTURE],
+            )
+        llm_dtype = self.language_model.config.dtype
+        assert isinstance(llm_dtype, torch.dtype)
+        self.llm_dtype = llm_dtype
+
+        with self._mark_tower_model(vllm_config, "audio"):
+            self.audio_tower = _build_qwen2_audio_encoder(config.audio_config).to(llm_dtype)
+            self.audio_projector = AudexProjector(
+                audio_encoder_hidden_size=config.audio_encoder_hidden_size,
+                intermediate_size=config.audio_projector_intermediate_size,
+                text_hidden_size=config.hidden_size,
+                activation=config.audio_projector_activation,
+                norm_eps=config.audio_projector_norm_eps,
+            ).to(llm_dtype)
+
+        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+
+    def get_language_model(self) -> nn.Module:
+        return self.language_model
+
+    def _parse_and_validate_audio_input(self, **kwargs: object):
+        input_features = kwargs.pop("input_features", None)
+        audio_num_clips = kwargs.pop("audio_num_clips", None)
+        if input_features is None:
+            return None
+        return AudexAudioFeatureInputs(
+            type="audio_features",
+            input_features=input_features,
+            audio_num_clips=audio_num_clips,
+            validate=False,
+        )
+
+    @torch.inference_mode()
+    def _encode_clips(self, input_features: torch.Tensor) -> torch.Tensor:
+        """(C, 128, 3000) -> (C, 750, hidden) via microbatched encoder+projector."""
+        param = next(self.audio_tower.parameters())
+        input_features = input_features.to(device=param.device, dtype=param.dtype)
+        out_chunks = []
+        for start in range(0, input_features.shape[0], AUDIO_ENCODER_MICRO_BATCH):
+            chunk = input_features[start : start + AUDIO_ENCODER_MICRO_BATCH]
+            hidden = self.audio_tower(input_features=chunk, return_dict=True).last_hidden_state
+            out_chunks.append(self.audio_projector(hidden))
+        return torch.cat(out_chunks, dim=0)
+
+    def _process_audio_input(self, audio_input) -> tuple[torch.Tensor, ...]:
+        input_features = audio_input["input_features"]
+        audio_num_clips = audio_input["audio_num_clips"]
+        if input_features.ndim == 4:
+            input_features = input_features.flatten(0, 1)
+
+        projected = self._encode_clips(input_features)
+        per_clip = projected.reshape(projected.shape[0], -1, projected.shape[-1])
+
+        if audio_num_clips is None:
+            num_clips_list = [projected.shape[0]]
+        elif isinstance(audio_num_clips, torch.Tensor):
+            num_clips_list = audio_num_clips.flatten().tolist()
+        else:
+            num_clips_list = [int(x) for x in audio_num_clips]
+
+        grouped: list[torch.Tensor] = []
+        offset = 0
+        for n in num_clips_list:
+            clips = per_clip[offset : offset + n]
+            grouped.append(clips.reshape(-1, clips.shape[-1]))
+            offset += n
+        return tuple(grouped)
+
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
+        audio_input = self._parse_and_validate_audio_input(**kwargs)
+        if audio_input is None:
+            return []
+        return self._process_audio_input(audio_input)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+        return self.language_model.model(input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds)
+
+    def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata=None) -> torch.Tensor | None:
+        return self.language_model.compute_logits(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Split the combined weight stream into LLM / audio_tower / projector.
+
+        vLLM streams *all* ``.safetensors`` in the model dir; route
+        ``audio_encoder.*``/``audio_projector.*`` to the audio modules and
+        leave the remaining ``model.*``/``lm_head.*`` names untouched so the
+        wrapped NemotronDenseForCausalLM loader applies its own qkv fusion
+        and remapping.
+        """
+        loaded: set[str] = set()
+        lm_weights: list[tuple[str, torch.Tensor]] = []
+        enc_sd: dict[str, torch.Tensor] = {}
+        proj_map: dict[str, torch.Tensor] = {}
+
+        for name, w in weights:
+            if name.startswith("audio_encoder."):
+                enc_sd[name[len("audio_encoder.") :]] = w
+            elif name == "audio_projector.norm.weight":
+                proj_map["norm_weight"] = w
+            elif name.startswith("audio_projector."):
+                proj_map[name[len("audio_projector.") :]] = w
+            else:
+                lm_weights.append((name, w))
+
+        lm_loaded = self.language_model.load_weights(lm_weights)
+        loaded |= {f"language_model.{n}" for n in lm_loaded}
+
+        if not enc_sd:
+            raise ValueError("No audio_encoder.* weights found in checkpoint stream.")
+        missing, unexpected = self.audio_tower.load_state_dict(enc_sd, strict=False)
+        if missing:
+            raise ValueError(f"Missing audio_tower params: {list(missing)[:8]}")
+        if unexpected:
+            raise ValueError(f"Unexpected audio_tower keys: {list(unexpected)[:8]}")
+        loaded |= {f"audio_tower.{k}" for k in enc_sd}
+
+        if set(proj_map) != {"norm_weight", "fc1.weight", "fc2.weight"}:
+            raise ValueError(f"Unexpected audio_projector keys: {sorted(proj_map)}")
+        self.audio_projector.load_state_dict(proj_map, strict=True)
+        loaded |= {f"audio_projector.{k}" for k in proj_map}
+
+        return loaded
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    AudexMultiModalProcessor,
+    info=AudexProcessingInfo,
+    dummy_inputs=AudexDummyInputsBuilder,
+)
+class NemotronHAudexForConditionalGeneration(
+    NemotronDenseAudexForConditionalGeneration,
+    IsHybrid,
+    HasInnerState,
+    SupportsMambaPrefixCaching,
+):
+    """Audex 30B-A3B audio-understanding model (``checkpoint_folder_full``).
+
+    Mirrors nvidia/Nemotron-Labs-Audex-30B-A3B's
+    ``inference_scripts_vllm/audioqa_scripts/audex_30b_a3b_vllm/modeling_audex_vllm.py``:
+    the same encoder/projector/processing as the dense 2B wrapper around
+    vLLM's native hybrid Mamba+MoE ``NemotronHForCausalLM``, whose state
+    hooks are delegated below.
+    """
+
+    _LM_ARCHITECTURE = "NemotronHForCausalLM"
+
+    is_hybrid = True
+    has_inner_state = True
+    supports_mamba_prefix_caching = True
+
+    @classmethod
+    def get_mamba_state_shape_from_config(cls, vllm_config):
+        from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
+
+        return NemotronHForCausalLM.get_mamba_state_shape_from_config(vllm_config)
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(cls, vllm_config):
+        from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
+
+        return NemotronHForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
+
+    @classmethod
+    def get_mamba_state_copy_func(cls):
+        from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
+
+        return NemotronHForCausalLM.get_mamba_state_copy_func()

vllm_omni/model_executor/models/audex/audex_omni_processing.py

@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0):
+#   inference_scripts_vllm/audioqa_scripts/audex_2b_vllm/audio_features.py
+#   inference_scripts_vllm/audioqa_scripts/audex_2b_vllm/processing_audex_vllm.py
+"""Multimodal processing for Audex audio understanding (S2T / audio QA).
+
+No unified HF ``ProcessorMixin`` exists for Audex (only a
+``WhisperFeatureExtractor`` under ``audio_preprocessor/``), so
+``_call_hf_processor`` is fully overridden: it tokenizes the prompt and
+builds NV-Whisper features with 30 s non-overlapping windows (last clip
+zero-padded), exactly mirroring the release HF path.
+
+Placeholder expansion uses vLLM's ``PromptReplacement``: a single
+``<so_embedding>`` per audio expands to ``<so_start>`` + N*``<so_embedding>``
++ ``<so_end>`` where N = num_clips * 750.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from transformers import BatchFeature
+from vllm.config.multimodal import BaseDummyOptions
+from vllm.inputs import MultiModalDataDict
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    PromptReplacement,
+    PromptUpdate,
+    PromptUpdateDetails,
+)
+
+SOUND_TOKEN = "<so_embedding>"
+SOUND_START_TOKEN = "<so_start>"
+SOUND_END_TOKEN = "<so_end>"
+
+SOUND_EMBEDDING_SIZE = 750
+SOUND_CLIP_DURATION = 30.0
+SOUND_TARGET_RATE = 16000
+
+MAX_AUDIO_SECONDS = 900
+MAX_AUDIO_CLIPS = 30
+MAX_AUDIO_TOKENS = MAX_AUDIO_CLIPS * SOUND_EMBEDDING_SIZE
+
+AUDIO_ENCODER_MICRO_BATCH = 8
+
+
+# ---------------------------------------------------------------- features
+
+
+def normalize_audio(audio: np.ndarray) -> np.ndarray:
+    """Return mono float32 audio in [-1, 1], matching the Megatron/HF path."""
+    audio = np.asarray(audio)
+    if audio.ndim == 2:
+        if audio.shape[1] <= 2:
+            audio = audio.mean(axis=1)
+        elif audio.shape[0] <= 2:
+            audio = audio.mean(axis=0)
+        else:
+            raise ValueError(f"Unsupported audio shape: {audio.shape}")
+
+    if audio.dtype == np.int16:
+        audio = audio.astype(np.float32) / 32768.0
+    elif audio.dtype != np.float32:
+        audio = audio.astype(np.float32)
+
+    max_abs = float(np.abs(audio).max()) if audio.size else 0.0
+    if max_abs > 1.0:
+        audio = audio / max_abs
+    return audio.astype(np.float32, copy=False)
+
+
+def split_audio_into_clips(
+    audio: np.ndarray,
+    sample_rate: int = SOUND_TARGET_RATE,
+    clip_duration: float = SOUND_CLIP_DURATION,
+) -> list[np.ndarray]:
+    """Split audio into fixed clips; pad the final clip to full length."""
+    audio = normalize_audio(audio)
+    clip_samples = int(round(sample_rate * clip_duration))
+    if clip_samples <= 0:
+        raise ValueError(f"Invalid clip_samples: {clip_samples}")
+    if audio.size == 0:
+        audio = np.zeros(1, dtype=np.float32)
+
+    num_clips = max(1, math.ceil(audio.shape[0] / clip_samples))
+    clips: list[np.ndarray] = []
+    for idx in range(num_clips):
+        start = idx * clip_samples
+        clip = audio[start : start + clip_samples]
+        if clip.shape[0] < clip_samples:
+            clip = np.pad(clip, (0, clip_samples - clip.shape[0]))
+        clips.append(clip.astype(np.float32, copy=False))
+    return clips
+
+
+def extract_whisper_features(
+    feature_extractor,
+    audio: np.ndarray,
+    sample_rate: int = SOUND_TARGET_RATE,
+    clip_duration: float = SOUND_CLIP_DURATION,
+) -> torch.Tensor:
+    """Return NV-Whisper input features shaped (num_clips, 128, 3000)."""
+    clips = split_audio_into_clips(audio, sample_rate=sample_rate, clip_duration=clip_duration)
+    features = feature_extractor(
+        clips,
+        sampling_rate=sample_rate,
+        return_tensors="pt",
+        padding="max_length",
+        return_attention_mask=False,
+    )
+    input_features = features.input_features
+    if input_features.ndim != 3:
+        raise ValueError(f"Expected 3D Whisper features, got {tuple(input_features.shape)}")
+    return input_features
+
+
+# ---------------------------------------------------------------- processor
+
+
+class AudexProcessingInfo(BaseProcessingInfo):
+    def get_hf_config(self):
+        return self.ctx.model_config.hf_config
+
+    def get_feature_extractor(self, **kwargs: object):
+        from transformers import AutoFeatureExtractor
+
+        config = self.get_hf_config()
+        path = getattr(config, "audio_preprocessor_path", None) or "audio_preprocessor"
+        candidate = Path(path)
+        if not candidate.is_absolute():
+            candidate = Path(self.ctx.model_config.model) / candidate
+        return AutoFeatureExtractor.from_pretrained(str(candidate))
+
+    def get_data_parser(self):
+        return MultiModalDataParser(target_sr=SOUND_TARGET_RATE, target_channels=1)
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": 1}
+
+    def get_mm_max_tokens_per_item(self, seq_len: int, mm_counts: Mapping[str, int] | None = None) -> Mapping[str, int]:
+        return {"audio": MAX_AUDIO_TOKENS}
+
+    def sound_token_ids(self) -> tuple[int, int, int]:
+        cfg = self.get_hf_config()
+        return (
+            int(cfg.sound_start_token_id),
+            int(cfg.sound_token_id),
+            int(cfg.sound_end_token_id),
+        )
+
+
+class AudexDummyInputsBuilder(BaseDummyInputsBuilder[AudexProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        return SOUND_TOKEN * mm_counts.get("audio", 0)
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+    ) -> MultiModalDataDict:
+        num_audios = mm_counts.get("audio", 0)
+        audio_len = int(SOUND_CLIP_DURATION * SOUND_TARGET_RATE)
+        overrides = mm_options.get("audio") if mm_options else None
+        return {"audio": self._get_dummy_audios(length=audio_len, num_audios=num_audios, overrides=overrides)}
+
+
+class AudexMultiModalProcessor(BaseMultiModalProcessor[AudexProcessingInfo]):
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, Any],
+        tok_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        tokenizer = self.info.get_tokenizer()
+        prompt_ids = tokenizer(prompt, add_special_tokens=False).input_ids
+
+        audios = mm_data.get("audios", []) or mm_data.get("audio", [])
+        if not audios:
+            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
+
+        feature_extractor = self.info.get_feature_extractor()
+        all_features: list[torch.Tensor] = []
+        audio_num_clips: list[int] = []
+        for item in audios:
+            audio = item[0] if isinstance(item, tuple) else item
+            audio = normalize_audio(np.asarray(audio))
+            duration = len(audio) / SOUND_TARGET_RATE
+            if duration > MAX_AUDIO_SECONDS:
+                raise ValueError(f"Audio duration {duration:.1f}s exceeds MAX_AUDIO_SECONDS={MAX_AUDIO_SECONDS}s.")
+            clips = split_audio_into_clips(audio, sample_rate=SOUND_TARGET_RATE, clip_duration=SOUND_CLIP_DURATION)
+            if len(clips) > MAX_AUDIO_CLIPS:
+                raise ValueError(f"Audio needs {len(clips)} clips > MAX_AUDIO_CLIPS={MAX_AUDIO_CLIPS}.")
+            feats = extract_whisper_features(
+                feature_extractor,
+                audio,
+                sample_rate=SOUND_TARGET_RATE,
+                clip_duration=SOUND_CLIP_DURATION,
+            )
+            all_features.append(feats)
+            audio_num_clips.append(feats.shape[0])
+
+        input_features = torch.cat(all_features, dim=0)
+        return BatchFeature(
+            dict(
+                input_ids=[prompt_ids],
+                input_features=input_features,
+                audio_num_clips=torch.tensor(audio_num_clips, dtype=torch.long),
+            ),
+            tensor_type="pt",
+        )
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        audio_num_clips = hf_inputs.get("audio_num_clips", torch.zeros(0, dtype=torch.long))
+        return dict(
+            input_features=MultiModalFieldConfig.flat_from_sizes("audio", audio_num_clips),
+            audio_num_clips=MultiModalFieldConfig.batched("audio"),
+        )
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, Any],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        so_start, so_embed, so_end = self.info.sound_token_ids()
+        out_mm_data = out_mm_kwargs.get_data()
+        num_clips = out_mm_data.get("audio_num_clips", torch.zeros(0, dtype=torch.long))
+        num_clips_list = num_clips.flatten().tolist() if hasattr(num_clips, "flatten") else list(num_clips)
+
+        def get_replacement_audex(item_idx: int):
+            n = int(num_clips_list[item_idx]) * SOUND_EMBEDDING_SIZE
+            if n < 1:
+                raise ValueError(f"Audio item {item_idx} produced 0 embeddings.")
+            full = [so_start] + [so_embed] * n + [so_end]
+            return PromptUpdateDetails.select_token_id(full, so_embed)
+
+        return [
+            PromptReplacement(
+                modality="audio",
+                target=[so_embed],
+                replacement=get_replacement_audex,
+            )
+        ]

vllm_omni/model_executor/models/audex/audex_thinker.py

@@ -0,0 +1,414 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from the NemotronDense vLLM plugin bundled with
+# nvidia/Nemotron-Labs-Audex-2B (nemotron_dense_vllm_plugin), which is not
+# installed as a package; the model code is carried here instead.
+"""Audex thinker: the Nemotron dense decoder-only LM that autoregressively
+emits ``<speechcodec_N>`` tokens for TTS.
+
+The thinker is a plain token-in/token-out vLLM model. Speech-codec token
+extraction happens in the stage input processors
+(``vllm_omni/model_executor/stage_input_processors/audex.py``), which read the
+sampled token stream directly, so no omni-specific output plumbing is needed
+here.
+"""
+
+from collections.abc import Iterable
+from itertools import islice
+from typing import Any
+
+import torch
+from torch import nn
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
+from vllm.sequence import IntermediateTensors
+
+
+class NemotronDenseRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class NemotronDenseMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.up_proj = ColumnParallelLinear(
+            input_size=hidden_size,
+            output_size=intermediate_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        self.act_fn = get_act_fn(hidden_act)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        up, _ = self.up_proj(x)
+        x = self.act_fn(up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class NemotronDenseAttention(nn.Module):
+    def __init__(
+        self,
+        config: Any,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        max_position_embeddings: int = 8192,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        cache_config: CacheConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = getattr(config, "head_dim", None)
+        if self.head_dim is None:
+            self.head_dim = self.hidden_size // self.total_num_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.max_position_embeddings = max_position_embeddings
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+        )
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class NemotronDenseDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: Any,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+        attention_bias = getattr(config, "attention_bias", False) or getattr(config, "bias", False)
+        self.self_attn = NemotronDenseAttention(
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads),
+            max_position_embeddings=max_position_embeddings,
+            quant_config=quant_config,
+            bias=attention_bias,
+            cache_config=cache_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = NemotronDenseMLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            bias=getattr(config, "mlp_bias", False),
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = NemotronDenseRMSNorm(config.hidden_size, eps=config.norm_eps)
+        self.post_attention_layernorm = NemotronDenseRMSNorm(config.hidden_size, eps=config.norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@support_torch_compile
+class NemotronDenseModel(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+
+        self.config = config
+        self.quant_config = quant_config
+
+        self.vocab_size = config.vocab_size
+
+        if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank):
+            self.embed_tokens = VocabParallelEmbedding(
+                self.vocab_size,
+                config.hidden_size,
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: NemotronDenseDecoderLayer(
+                config=config,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=prefix,
+            ),
+            prefix=f"{prefix}.layers",
+        )
+        if get_pp_group().is_last_rank:
+            self.norm = NemotronDenseRMSNorm(config.hidden_size, eps=config.norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states"], config.hidden_size
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_input_ids(input_ids)
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+
+        for layer in islice(self.layers, self.start_layer, self.end_layer):
+            hidden_states = layer(positions, hidden_states)
+
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors({"hidden_states": hidden_states})
+
+        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
+                param = params_dict[scale_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
+                weight_loader(param, loaded_weight)
+                loaded_params.add(scale_name)
+                continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if is_pp_missing_parameter(name, self):
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+
+                break
+            else:
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+
+                if is_pp_missing_parameter(name, self):
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class NemotronDenseForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+    }
+
+    embedding_modules = {
+        "embed_tokens": "input_embeddings",
+        "lm_head": "output_embeddings",
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+
+        assert hasattr(config, "norm_eps") and hasattr(config, "num_hidden_layers")
+
+        self.config = config
+        self.quant_config = quant_config
+
+        self.model = NemotronDenseModel(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+            if config.tie_word_embeddings:
+                self.lm_head.weight = self.model.embed_tokens.weight
+
+            logit_scale = getattr(config, "logit_scale", 1.0)
+            self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale)
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | IntermediateTensors:
+        # The omni AR runner forwards extra bookkeeping kwargs (e.g.
+        # sampling_metadata); this plain LM ignores them.
+        return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        sampling_metadata: Any = None,
+    ) -> torch.Tensor | None:
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/model_executor/models/audex/audex_xcodec.py

@@ -0,0 +1,178 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Audex TTA stage 1: XCodec1 RVQ codes → waveform.
+
+Unlike the streaming causal speech decoder used for TTS, XCodec1 is a CNN
+codec that the official flow decodes over the full sequence, so this stage
+runs sync full-payload only: each request delivers its complete
+de-interleaved ``[frames, 4]`` codec-id payload once at finish and the
+decode happens in a single call.
+
+The checkpoint (``hf-audio/xcodec-hubert-general-balanced``) is external to
+the Audex repo and loads through transformers remote code; the model
+skeleton is built from config at startup for vLLM's memory profiler and the
+weights arrive via the standard ``load_weights`` iteration over the
+snapshot's safetensors.
+"""
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+from torch import nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.audex.tta import (
+    XCODEC1_CODEBOOK_SIZE,
+    XCODEC1_NUM_CODEBOOKS,
+)
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+logger = init_logger(__name__)
+
+# Official cap on decoded frames (~10 s of audio at the codec frame rate).
+_DEFAULT_MAX_TTA_FRAMES = 500
+
+
+def _meta_str(value: Any) -> str | None:
+    if isinstance(value, list | tuple):
+        return _meta_str(value[0]) if value else None
+    if isinstance(value, str):
+        return value
+    return None
+
+
+def _codes_from_runtime_info(runtime_info: dict[str, Any] | None) -> torch.Tensor | None:
+    """Extract the ``[frames, 4]`` codec-id payload from the connector info."""
+    if not isinstance(runtime_info, dict):
+        return None
+    codes = runtime_info.get("codes")
+    if not isinstance(codes, dict):
+        return None
+    audio = codes.get("audio")
+    if isinstance(audio, torch.Tensor):
+        return audio
+    if isinstance(audio, list | tuple) and audio:
+        return torch.as_tensor(audio, dtype=torch.long)
+    return None
+
+
+class AudexXCodec1(nn.Module):
+    """Stage-1 model for Audex TTA (GenerationModelRunner)."""
+
+    input_modalities = "audio"
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.model_path = vllm_config.model_config.model
+
+        self.have_multimodal_outputs = True
+        self.has_preprocess = False
+        self.has_postprocess = False
+        self.enable_update_additional_information = True
+        self.requires_raw_input_tokens = True
+
+        from transformers import AutoConfig, AutoModel
+
+        config = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True)
+        self.codec = AutoModel.from_config(config, trust_remote_code=True)
+        self.codec.eval()
+        self._sample_rate = int(getattr(config, "sampling_rate", 16000))
+        self._max_frames = int(getattr(config, "max_tta_frames", _DEFAULT_MAX_TTA_FRAMES))
+
+    # -------------------- vLLM model interface --------------------
+
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        # This stage ignores token embeddings; keep a stable dummy for the runner.
+        if input_ids.numel() == 0:
+            return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
+        return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
+
+    def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None:
+        return None
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        state_dict = dict(weights)
+        self.codec.load_state_dict(state_dict, strict=True)
+        device = self.vllm_config.device_config.device
+        self.codec.to(device=device, dtype=torch.float32)
+        return {f"codec.{name}" for name in state_dict}
+
+    def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput:
+        if isinstance(model_outputs, OmniOutput):
+            return model_outputs
+        # The runner may unpack the NamedTuple in transit; rebuild it.
+        if isinstance(model_outputs, tuple) and len(model_outputs) == len(OmniOutput._fields):
+            return OmniOutput(*model_outputs)
+        raise TypeError(f"AudexXCodec1 expected OmniOutput, got {type(model_outputs)}")
+
+    # -------------------- decoding --------------------
+
+    def _decode_codes(self, req_id: str | None, codes: torch.Tensor) -> torch.Tensor:
+        """Decode a ``[frames, 4]`` codec-id payload to a mono waveform."""
+        if codes.numel() == 0:
+            # Zero codec tokens for the whole request. Do NOT raise: an
+            # exception in forward() kills the engine core for every request.
+            # Returning empty audio lets the serving layer fail this request.
+            logger.error("Audex TTA: request %s finished with no codec tokens; no audio generated", req_id)
+            return torch.empty(0, dtype=torch.float32)
+        if codes.ndim != 2 or codes.shape[1] != XCODEC1_NUM_CODEBOOKS:
+            logger.error(
+                "Audex TTA: request %s payload has shape %s, expected [frames, %d]; no audio generated",
+                req_id,
+                tuple(codes.shape),
+                XCODEC1_NUM_CODEBOOKS,
+            )
+            return torch.empty(0, dtype=torch.float32)
+
+        device = self.vllm_config.device_config.device
+        with torch.inference_mode():
+            frames = codes[: self._max_frames].to(device=device, dtype=torch.long)
+            # [frames, 4] codec ids -> [1, 4, frames] per-layer indices.
+            layer_offsets = (
+                torch.arange(XCODEC1_NUM_CODEBOOKS, device=frames.device, dtype=frames.dtype) * XCODEC1_CODEBOOK_SIZE
+            )
+            per_layer = (frames - layer_offsets.view(1, -1)).clamp(0, XCODEC1_CODEBOOK_SIZE - 1)
+            stacked = per_layer.transpose(0, 1).unsqueeze(0)
+            decoded = self.codec.decode(stacked)
+            audio = decoded.audio_values if hasattr(decoded, "audio_values") else decoded[0]
+            return audio.reshape(-1).to(dtype=torch.float32, device="cpu")
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: Any = None,
+        inputs_embeds: torch.Tensor | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        runtime_infos = runtime_additional_information or []
+        seq_token_counts = kwargs.get("seq_token_counts")
+        if seq_token_counts is not None:
+            num_reqs = len(seq_token_counts)
+        elif input_ids is not None and input_ids.numel() > 0 and not runtime_infos:
+            num_reqs = 1
+        else:
+            num_reqs = len(runtime_infos)
+
+        audios: list[torch.Tensor] = []
+        sr_tensor = torch.tensor(self._sample_rate, dtype=torch.int32)
+        for i in range(num_reqs):
+            info = runtime_infos[i] if i < len(runtime_infos) else {}
+            info = info if isinstance(info, dict) else {}
+            meta = info.get("meta", {})
+            meta = meta if isinstance(meta, dict) else {}
+            req_id = _meta_str(meta.get("req_id"))
+            codes = _codes_from_runtime_info(info)
+            if codes is None:
+                audios.append(torch.empty(0, dtype=torch.float32))
+                continue
+            audios.append(self._decode_codes(req_id, codes))
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"model_outputs": audios, "sr": [sr_tensor] * num_reqs},
+        )

vllm_omni/model_executor/models/audex/cfg.py

@@ -0,0 +1,567 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0):
+#   inference_scripts_vllm/audiogen_scripts/cfg_logits_processor.py
+#   inference_scripts_vllm/audiogen_scripts/vllm_cfg_patch.py
+# Divergences from the official files are commented at the divergence site.
+"""Classifier-free guidance for the Audex thinker stage.
+
+CFG pairs a conditional request with an unconditional (null-prompt)
+companion in the same engine and blends their logits every step:
+
+    blended = uncond + cfg_scale * (cond - uncond)
+
+Both rows receive the blended logits and the sampled token is copied from
+the cond row to the uncond row, so the two sequences stay token-identical.
+Requests opt in via ``SamplingParams.extra_args``:
+
+    cond:   {"cfg_scale": 1.5, "cfg_role": "cond",   "cfg_pair_id": <id>}
+    uncond: {"cfg_scale": 1.5, "cfg_role": "uncond", "cfg_pair_id": <id>}
+
+Blending is only correct when both pair members decode the same position in
+the same engine step, so ``apply_cfg_patches()`` makes the vLLM scheduler
+pair-aware: a lone member waits for its partner, partners stay adjacent,
+their ``num_computed_tokens`` are equalized after every schedule, and they
+finish together. The patch is a no-op for engines that never see a CFG
+request (the pair registry stays empty).
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from typing import Any
+
+import torch
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.sampling_params import SamplingParams
+from vllm.v1.sample.logits_processor import (
+    BatchUpdate,
+    LogitsProcessor,
+    MoveDirectionality,
+)
+
+logger = init_logger(__name__)
+
+_COND = "cond"
+_UNCOND = "uncond"
+
+
+def _cfg_extra(request: Any, key: str) -> Any:
+    sampling_params = getattr(request, "sampling_params", None)
+    extra_args = getattr(sampling_params, "extra_args", None)
+    if extra_args:
+        return extra_args.get(key)
+    return None
+
+
+class AudexCFGLogitsProcessor(LogitsProcessor):
+    """Blend cond/uncond logits for classifier-free guidance.
+
+    Pairs are matched by ``cfg_pair_id``. For each pair the blended logits
+    are written to *both* rows so the sampler picks the same token; the
+    post-sampling copy of the cond token into the uncond slot is installed
+    by :meth:`_ensure_sample_patched` (which must run inside each worker
+    process, hence from ``__init__`` rather than a main-process patch).
+    """
+
+    _sample_patched = False
+
+    @classmethod
+    def validate_params(cls, params: SamplingParams) -> None:
+        extra_args = params.extra_args
+        if not extra_args:
+            return
+        role = extra_args.get("cfg_role")
+        if role is not None and role not in (_COND, _UNCOND):
+            raise ValueError(f"cfg_role must be 'cond' or 'uncond', got '{role}'")
+        scale = extra_args.get("cfg_scale")
+        if scale is not None and (not isinstance(scale, int | float) or scale < 1.0):
+            raise ValueError(f"cfg_scale must be >= 1.0, got {scale}")
+
+    def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool) -> None:
+        self._info: dict[int, dict[str, Any]] = {}
+        self._output_tokens: dict[int, list[int]] = {}
+        self._pairs: list[tuple[int, int, float]] = []
+        self._dirty = True
+        self._ensure_sample_patched()
+
+    @classmethod
+    def _ensure_sample_patched(cls) -> None:
+        """Install the post-sampling cond→uncond token copy (once per process)."""
+        if cls._sample_patched:
+            return
+        cls._sample_patched = True
+
+        # Divergence from the official file: it patches vLLM's
+        # ``GPUModelRunner._sample``, but vllm-omni's AR runner overrides
+        # ``_sample``, so the base-class patch would never run for the Audex
+        # thinker stage. Patch the omni AR runner instead.
+        from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+        orig_sample = GPUARModelRunner._sample
+
+        def _sample_with_cfg_sync(self, logits, spec_decode_metadata):
+            sampler_output = orig_sample(self, logits, spec_decode_metadata)
+            for proc in self.input_batch.logitsprocs.all:
+                if isinstance(proc, AudexCFGLogitsProcessor) and proc._pairs:
+                    sampled = sampler_output.sampled_token_ids
+                    num_rows = sampled.shape[0] if hasattr(sampled, "shape") else len(sampled)
+                    for cond_idx, uncond_idx, _ in proc._pairs:
+                        # Steps that sample fewer rows than the persistent
+                        # batch (pure-prefill / partially scheduled steps)
+                        # have nothing to sync for this pair; the scheduler
+                        # pair patches keep both members step-locked on the
+                        # decode steps that matter.
+                        if cond_idx < num_rows and uncond_idx < num_rows:
+                            sampled[uncond_idx] = sampled[cond_idx]
+                    break
+            return sampler_output
+
+        _sample_with_cfg_sync._audex_cfg_sync = True  # type: ignore[attr-defined]
+        GPUARModelRunner._sample = _sample_with_cfg_sync
+        logger.info("AudexCFGLogitsProcessor: patched GPUARModelRunner._sample for CFG token sync")
+
+    def is_argmax_invariant(self) -> bool:
+        return False
+
+    def _reset(self) -> None:
+        self._info.clear()
+        self._output_tokens.clear()
+        self._pairs.clear()
+        self._dirty = True
+
+    def update_state(self, batch_update: BatchUpdate | None) -> None:
+        if batch_update is None:
+            return
+
+        for idx in batch_update.removed:
+            self._info.pop(idx, None)
+            self._output_tokens.pop(idx, None)
+
+        if not self._info and batch_update.added:
+            self._reset()
+
+        for idx, params, _, output_token_ids in batch_update.added:
+            extra_args = params.extra_args if params else None
+            if extra_args and extra_args.get("cfg_role") in (_COND, _UNCOND):
+                self._info[idx] = {
+                    "role": extra_args["cfg_role"],
+                    "cfg_scale": float(extra_args.get("cfg_scale", 1.0)),
+                    "pair_id": extra_args.get("cfg_pair_id"),
+                }
+                self._output_tokens[idx] = output_token_ids
+            else:
+                self._info.pop(idx, None)
+                self._output_tokens.pop(idx, None)
+        self._dirty = True
+
+        if self._info:
+            for src, dst, direction in batch_update.moved:
+                src_info = self._info.pop(src, None)
+                dst_info = self._info.pop(dst, None)
+                src_tokens = self._output_tokens.pop(src, None)
+                dst_tokens = self._output_tokens.pop(dst, None)
+                if src_info is not None:
+                    self._info[dst] = src_info
+                if src_tokens is not None:
+                    self._output_tokens[dst] = src_tokens
+                if direction == MoveDirectionality.SWAP:
+                    if dst_info is not None:
+                        self._info[src] = dst_info
+                    if dst_tokens is not None:
+                        self._output_tokens[src] = dst_tokens
+            self._dirty = True
+
+    def _rebuild_pairs(self) -> None:
+        by_pair: dict[str, dict[str, tuple[int, float]]] = {}
+        for idx, info in self._info.items():
+            pair_id = info.get("pair_id")
+            if pair_id is None:
+                continue
+            by_pair.setdefault(pair_id, {})[info["role"]] = (idx, info["cfg_scale"])
+
+        self._pairs = [
+            (roles[_COND][0], roles[_UNCOND][0], roles[_COND][1])
+            for roles in by_pair.values()
+            if _COND in roles and _UNCOND in roles
+        ]
+        self._dirty = False
+
+    def apply(self, logits: torch.Tensor) -> torch.Tensor:
+        if not self._info:
+            return logits
+
+        if self._dirty:
+            self._rebuild_pairs()
+
+        num_rows = logits.shape[0]
+        for cond_idx, uncond_idx, cfg_scale in self._pairs:
+            # Rows beyond this step's logits (pure-prefill / partially
+            # scheduled steps) have nothing to blend; decode steps carry
+            # both pair rows thanks to the scheduler pair patches.
+            if cond_idx >= num_rows or uncond_idx >= num_rows:
+                continue
+            blended = logits[uncond_idx] + cfg_scale * (logits[cond_idx] - logits[uncond_idx])
+            logits[cond_idx] = blended
+            logits[uncond_idx] = blended
+
+        return logits
+
+
+# ---------------------------------------------------------------------------
+# Scheduler pair synchronization
+# ---------------------------------------------------------------------------
+
+
+def _wait_queues(scheduler: Any) -> list[Any]:
+    # Divergence from the official (vLLM 0.20-era) file: vLLM 0.24 parks
+    # structurally not-ready requests in a second ``skipped_waiting`` queue;
+    # pair-hold must cover both queues or a lone member parked there would
+    # be admitted without its partner.
+    queues = [scheduler.waiting]
+    skipped = getattr(scheduler, "skipped_waiting", None)
+    if skipped is not None:
+        queues.append(skipped)
+    return queues
+
+
+def _pair_complete(scheduler: Any, request_id: str, blocked_ids: set[str] | None = None) -> bool:
+    pair_id = scheduler._cfg_req_to_pair.get(request_id)
+    if pair_id is None:
+        return True
+    roles = scheduler._cfg_pairs.get(pair_id, {})
+    if len(roles) != 2 or not all(rid in scheduler.requests for rid in roles.values()):
+        return False
+    if blocked_ids:
+        # A partner parked in skipped_waiting (e.g. waiting for remote KVs)
+        # may fail promotion this step while this member gets scheduled from
+        # waiting — hold this member until the partner is schedulable.
+        for rid in roles.values():
+            if rid != request_id and rid in blocked_ids:
+                return False
+    return True
+
+
+# Schedule steps a lone pair member may wait for its partner. Partners arrive
+# within a handful of steps in practice; exhausting this budget means the
+# companion was never created (e.g. prompt expansion failed), and the request
+# is released to decode unguided instead of hanging forever.
+_MAX_PAIR_HOLD_STEPS = 512
+
+
+def _drop_stale_pair(scheduler: Any, request_id: str) -> None:
+    pair_id = scheduler._cfg_req_to_pair.get(request_id)
+    if pair_id is None:
+        return
+    for rid in scheduler._cfg_pairs.pop(pair_id, {}).values():
+        scheduler._cfg_req_to_pair.pop(rid, None)
+    logger.warning(
+        "CFG pair %s: partner of %s never arrived after %d schedule steps; releasing the request unguided",
+        pair_id,
+        request_id,
+        _MAX_PAIR_HOLD_STEPS,
+    )
+
+
+def _hold_incomplete_pairs(scheduler: Any) -> list[tuple[Any, Any]]:
+    """Pull CFG requests whose partner is not yet admitted out of the wait queues.
+
+    The engine schedules as soon as the first member of a pair arrives over
+    IPC; without the hold, the lone member prefills one step early and the
+    pair is permanently offset by one token. Held requests are re-prepended
+    after the schedule step.
+    """
+    hold_counts: dict[str, int] = getattr(scheduler, "_cfg_hold_counts", None) or {}
+    scheduler._cfg_hold_counts = hold_counts
+
+    skipped = getattr(scheduler, "skipped_waiting", None)
+    blocked_ids = {req.request_id for req in list(skipped)} if skipped is not None else set()
+
+    held: list[tuple[Any, Any]] = []
+    for queue in _wait_queues(scheduler):
+        # A member is only "blocked" for its partner's sake when it sits in
+        # the OTHER queue; members of the queue being scanned move together.
+        partner_blockers = blocked_ids if queue is not skipped else set()
+        held_here = []
+        for request in list(queue):
+            if _pair_complete(scheduler, request.request_id, partner_blockers):
+                hold_counts.pop(request.request_id, None)
+                continue
+            count = hold_counts.get(request.request_id, 0) + 1
+            if count > _MAX_PAIR_HOLD_STEPS:
+                hold_counts.pop(request.request_id, None)
+                _drop_stale_pair(scheduler, request.request_id)
+                continue
+            hold_counts[request.request_id] = count
+            held_here.append(request)
+        if held_here:
+            queue.remove_requests(held_here)
+            held.extend((queue, req) for req in held_here)
+    return held
+
+
+def _release_held(held: list[tuple[Any, Any]]) -> None:
+    for queue, request in reversed(held):
+        queue.prepend_request(request)
+
+
+def _reorder_waiting_for_cfg(scheduler: Any) -> None:
+    """Move CFG pair partners adjacent in the FCFS waiting queue."""
+    waiting = scheduler.waiting
+    # Only the FCFS queue (a deque) has caller-controlled ordering; priority
+    # queues order themselves and pair sync then relies on hold + equalize.
+    if not isinstance(waiting, deque) or len(waiting) < 2:
+        return
+
+    requests = list(waiting)
+    waiting.clear()
+
+    seen: set[str] = set()
+    result: list[Any] = []
+    for request in requests:
+        rid = request.request_id
+        if rid in seen:
+            continue
+        seen.add(rid)
+        result.append(request)
+
+        pair_id = scheduler._cfg_req_to_pair.get(rid)
+        if pair_id is None:
+            continue
+        for partner_id in scheduler._cfg_pairs.get(pair_id, {}).values():
+            if partner_id != rid and partner_id not in seen:
+                for candidate in requests:
+                    if candidate.request_id == partner_id:
+                        seen.add(partner_id)
+                        result.append(candidate)
+                        break
+
+    waiting.extend(result)
+
+
+def _equalize_cfg_pair_progress(scheduler: Any, scheduler_output: Any) -> None:
+    """Bring both members of every CFG pair to the same ``num_computed_tokens``.
+
+    Prefix-cache hits or unequal chunked-prefill budget can leave one member
+    ahead after a schedule step; reduce the faster member's allocation so
+    both land on the same position. Over-allocated KV blocks remain and are
+    consumed next step — nothing leaks.
+    """
+    for roles in scheduler._cfg_pairs.values():
+        cond_id = roles.get(_COND)
+        uncond_id = roles.get(_UNCOND)
+        if not cond_id or not uncond_id:
+            continue
+
+        cond_sched = scheduler_output.num_scheduled_tokens.get(cond_id, 0)
+        uncond_sched = scheduler_output.num_scheduled_tokens.get(uncond_id, 0)
+        if cond_sched == 0 or uncond_sched == 0:
+            continue
+
+        cond_req = scheduler.requests.get(cond_id)
+        uncond_req = scheduler.requests.get(uncond_id)
+        if not cond_req or not uncond_req:
+            continue
+
+        if cond_req.num_computed_tokens == uncond_req.num_computed_tokens:
+            continue
+
+        target = min(cond_req.num_computed_tokens, uncond_req.num_computed_tokens)
+
+        feasible = all(
+            req.num_computed_tokens - target < sched
+            for req, sched in ((cond_req, cond_sched), (uncond_req, uncond_sched))
+        )
+        if not feasible:
+            continue
+
+        for req_id, req, orig_sched in (
+            (cond_id, cond_req, cond_sched),
+            (uncond_id, uncond_req, uncond_sched),
+        ):
+            diff = req.num_computed_tokens - target
+            if diff > 0:
+                req.num_computed_tokens = target
+                scheduler_output.num_scheduled_tokens[req_id] = orig_sched - diff
+                scheduler_output.total_num_scheduled_tokens -= diff
+                logger.debug(
+                    "CFG equalize %s: %d -> %d scheduled, computed -> %d",
+                    req_id,
+                    orig_sched,
+                    orig_sched - diff,
+                    target,
+                )
+
+
+def _drop_split_pairs(scheduler: Any) -> None:
+    """Deregister a desynced CFG pair (one member preempted / left behind).
+
+    Manual re-preemption surgery is unsafe against the upstream scheduler.
+    Instead, a pair that split anyway (KV-pressure preemption of exactly one
+    member, or unrecoverable progress skew) is dropped from the registry:
+    the cond request continues UNGUIDED (its logits are never blended again)
+    rather than being re-paired against a desynced partner, and the
+    companion's discarded output no longer matters.
+    """
+    if not scheduler._cfg_pairs:
+        return
+
+    running_ids = {req.request_id for req in scheduler.running}
+    to_drop: list[tuple[str, str, str, str]] = []
+    for pair_id, roles in scheduler._cfg_pairs.items():
+        cond_id = roles.get(_COND)
+        uncond_id = roles.get(_UNCOND)
+        if cond_id is None or uncond_id is None:
+            continue
+        if cond_id not in scheduler.requests or uncond_id not in scheduler.requests:
+            continue
+
+        cond_running = cond_id in running_ids
+        uncond_running = uncond_id in running_ids
+        if cond_running != uncond_running:
+            # Split only counts once decoding started: a lone member that has
+            # computed tokens while its partner was pushed back is desynced.
+            runner = scheduler.requests.get(cond_id if cond_running else uncond_id)
+            if runner is not None and runner.num_computed_tokens > 0:
+                to_drop.append((pair_id, cond_id, uncond_id, "one member preempted"))
+            continue
+
+        if cond_running and uncond_running:
+            cond_req = scheduler.requests.get(cond_id)
+            uncond_req = scheduler.requests.get(uncond_id)
+            if cond_req and uncond_req and cond_req.num_computed_tokens != uncond_req.num_computed_tokens:
+                logger.warning(
+                    "CFG pair progress mismatch (equalize will retry): %s@%d vs %s@%d",
+                    cond_id,
+                    cond_req.num_computed_tokens,
+                    uncond_id,
+                    uncond_req.num_computed_tokens,
+                )
+
+    for pair_id, cond_id, uncond_id, reason in to_drop:
+        scheduler._cfg_pairs.pop(pair_id, None)
+        scheduler._cfg_req_to_pair.pop(cond_id, None)
+        scheduler._cfg_req_to_pair.pop(uncond_id, None)
+        logger.warning(
+            "CFG pair %s split (%s); releasing cond=%s unguided, uncond=%s output is discarded",
+            pair_id,
+            reason,
+            cond_id,
+            uncond_id,
+        )
+
+
+_patches_applied = False
+
+
+def cfg_patches_applied() -> bool:
+    """True once :func:`apply_cfg_patches` has patched the scheduler in this process."""
+    return _patches_applied
+
+
+def _patch_scheduler() -> None:
+    from vllm.v1.core.sched.scheduler import Scheduler
+
+    orig_init = Scheduler.__init__
+    orig_add = Scheduler.add_request
+    orig_finish = Scheduler.finish_requests
+    orig_schedule = Scheduler.schedule
+
+    def _init(self, *args: Any, **kwargs: Any) -> None:
+        orig_init(self, *args, **kwargs)
+        self._cfg_pairs: dict[str, dict[str, str]] = {}
+        self._cfg_req_to_pair: dict[str, str] = {}
+
+    def _add(self, request: Any) -> None:
+        orig_add(self, request)
+        pair_id = _cfg_extra(request, "cfg_pair_id")
+        role = _cfg_extra(request, "cfg_role")
+        if pair_id and role:
+            self._cfg_pairs.setdefault(pair_id, {})[role] = request.request_id
+            self._cfg_req_to_pair[request.request_id] = pair_id
+            # An asymmetric prefix-cache hit (the cond prompt shares cached
+            # branches; the <unk> null prompt does not) can advance one
+            # member past what post-schedule equalization can claw back.
+            # Skip prefix-cache reads for CFG requests entirely.
+            if hasattr(request, "skip_reading_prefix_cache"):
+                request.skip_reading_prefix_cache = True
+
+    def _finish(self, request_ids: Any, finished_status: Any) -> Any:
+        if request_ids is None:
+            result = orig_finish(self, request_ids, finished_status)
+            self._cfg_pairs.clear()
+            self._cfg_req_to_pair.clear()
+            return result
+
+        if isinstance(request_ids, str):
+            request_ids = {request_ids}
+        else:
+            request_ids = set(request_ids)
+
+        # Finish both members together: a surviving partner would decode
+        # unpaired garbage (its logits are never blended again).
+        partner_ids: set[str] = set()
+        for req_id in request_ids:
+            pair_id = self._cfg_req_to_pair.get(req_id)
+            if pair_id is None:
+                continue
+            for rid in self._cfg_pairs.get(pair_id, {}).values():
+                if rid != req_id and rid in self.requests:
+                    partner_ids.add(rid)
+
+        all_ids = request_ids | partner_ids
+        result = orig_finish(self, all_ids, finished_status)
+
+        for req_id in all_ids:
+            pair_id = self._cfg_req_to_pair.pop(req_id, None)
+            if pair_id:
+                self._cfg_pairs.pop(pair_id, None)
+        return result
+
+    def _schedule(self, *args: Any, **kwargs: Any) -> Any:
+        if not self._cfg_pairs:
+            return orig_schedule(self, *args, **kwargs)
+
+        _reorder_waiting_for_cfg(self)
+        held = _hold_incomplete_pairs(self)
+
+        # Chunked prefill can split one pair member's prompt across steps
+        # while the other fits in a single step; halving the long-prefill
+        # threshold makes both chunk, and equalize aligns the remainder.
+        scheduler_config = getattr(self, "scheduler_config", None)
+        orig_threshold = getattr(scheduler_config, "long_prefill_token_threshold", None)
+        if scheduler_config is not None and orig_threshold is not None:
+            scheduler_config.long_prefill_token_threshold = self.max_num_scheduled_tokens // 2
+        try:
+            scheduler_output = orig_schedule(self, *args, **kwargs)
+        finally:
+            if scheduler_config is not None and orig_threshold is not None:
+                scheduler_config.long_prefill_token_threshold = orig_threshold
+            _release_held(held)
+
+        _equalize_cfg_pair_progress(self, scheduler_output)
+        _drop_split_pairs(self)
+        return scheduler_output
+
+    _schedule._audex_cfg_patched = True  # type: ignore[attr-defined]
+
+    Scheduler.__init__ = _init
+    Scheduler.add_request = _add
+    Scheduler.finish_requests = _finish
+    Scheduler.schedule = _schedule
+
+
+def apply_cfg_patches() -> None:
+    """Make the vLLM v1 scheduler CFG-pair-aware. Idempotent, per process.
+
+    Must run in the engine-core process before ``Scheduler`` is constructed
+    (the ``__init__`` wrapper installs the pair registry).
+    """
+    global _patches_applied
+    if _patches_applied:
+        return
+    _patches_applied = True
+    logger.info("Applying Audex CFG scheduler patches to vLLM v1")
+    _patch_scheduler()

vllm_omni/model_executor/models/audex/checkpoint.py

@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Audex checkpoint layout preparation.
+
+The nvidia/Nemotron-Labs-Audex-2B repo deduplicates weights: the
+``checkpoint_folder_audiogen`` (and ``checkpoint_folder_textonly``) index
+files reference safetensors shards that physically live only under
+``checkpoint_folder_full``. The official repo ships a prepare script that
+symlinks them; this replicates that step automatically so users can pass the
+repo root without any manual preparation.
+"""
+
+import json
+import os
+import shutil
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_WEIGHT_SOURCE_FOLDER = "checkpoint_folder_full"
+
+# Download profiles: each deployment pulls only the repo subset its stages
+# need. The dedup weight shards are appended per-model at call time from
+# the audiogen index (see _dedup_shard_patterns); the static lists here
+# hold only the model-independent folders.
+_SNAPSHOT_PROFILE_PATTERNS = {
+    "tts": [
+        "config.json",
+        "checkpoint_folder_audiogen/*",
+        "audex_causal_speech_decoder/*",
+    ],
+    # TTA reuses the audiogen thinker; waveform decode happens in the
+    # external XCodec1 checkpoint (see ensure_xcodec1_snapshot), so the
+    # bundled speech decoder is not pulled.
+    "tta": [
+        "config.json",
+        "checkpoint_folder_audiogen/*",
+    ],
+    # The audio-capable full checkpoint: both weight shards, audio
+    # preprocessor assets, chat template, and remote-code modeling files.
+    "full": [
+        "config.json",
+        f"{_WEIGHT_SOURCE_FOLDER}/*",
+    ],
+}
+
+XCODEC1_DEFAULT_REPO = "hf-audio/xcodec-hubert-general-balanced"
+
+_AUDIOGEN_INDEX_FILE = "checkpoint_folder_audiogen/model.safetensors.index.json"
+
+
+def _download_audiogen_index(model: str) -> str:
+    """Resolve the audiogen dedup index to a local path (tiny, cache-first).
+
+    Cache-first keeps warm-cache boots offline-safe and avoids a Hub
+    metadata round trip per stage-engine boot; a cache miss falls back to
+    the network.
+    """
+    from huggingface_hub import hf_hub_download
+
+    try:
+        return hf_hub_download(model, _AUDIOGEN_INDEX_FILE, local_files_only=True)
+    except Exception:
+        return hf_hub_download(model, _AUDIOGEN_INDEX_FILE)
+
+
+def _dedup_shard_patterns(model: str) -> list[str]:
+    """Shards under the full checkpoint referenced by the audiogen index.
+
+    The audiogen folder deduplicates its LM weights into
+    ``checkpoint_folder_full`` shards (1 of 2 on the 2B, a subset of 14 on
+    the 30B-A3B); the ``tts``/``tta`` profiles must download exactly those.
+    Falls back to every full-checkpoint shard when the index is unreadable.
+    """
+    try:
+        with open(_download_audiogen_index(model)) as f:
+            weight_map = json.load(f).get("weight_map", {})
+        shards = sorted(set(weight_map.values()))
+        if shards:
+            return [f"{_WEIGHT_SOURCE_FOLDER}/{shard}" for shard in shards]
+    except Exception as exc:
+        logger.warning(
+            "Could not resolve the audiogen dedup shards for %s from %s (%s); downloading all full-checkpoint shards",
+            model,
+            _AUDIOGEN_INDEX_FILE,
+            exc,
+        )
+    return [f"{_WEIGHT_SOURCE_FOLDER}/model-*.safetensors"]
+
+
+def ensure_audex_snapshot(model: str, profile: str = "tts") -> str:
+    """Resolve ``model`` to a local repo-root directory, downloading if needed.
+
+    Local paths pass through untouched. For an HF repo id, download (or reuse
+    from cache) exactly the subset the requesting stage's ``profile`` needs,
+    so per-stage ``model_subdir`` joins land on a real snapshot instead of
+    the repo-id string on a fresh cache. Falls back to a cached snapshot
+    when offline.
+    """
+    if os.path.isdir(model):
+        return model
+
+    patterns = _SNAPSHOT_PROFILE_PATTERNS.get(profile)
+    if patterns is None:
+        raise ValueError(
+            f"Unknown Audex snapshot profile {profile!r}; expected one of {sorted(_SNAPSHOT_PROFILE_PATTERNS)}"
+        )
+
+    patterns = list(patterns)
+    if profile in ("tts", "tta"):
+        patterns += _dedup_shard_patterns(model)
+
+    from huggingface_hub import snapshot_download
+
+    try:
+        return snapshot_download(model, allow_patterns=patterns)
+    except Exception as download_exc:
+        try:
+            return snapshot_download(model, allow_patterns=patterns, local_files_only=True)
+        except Exception:
+            raise RuntimeError(
+                f"Could not resolve the Audex repo {model!r} (profile {profile!r}): the download "
+                "failed and no cached snapshot with the required folders was found."
+            ) from download_exc
+
+
+def ensure_xcodec1_snapshot(model: str | None) -> str:
+    """Resolve the external XCodec1 checkpoint (TTA waveform decoder).
+
+    Accepts a local directory or an HF repo id; ``None``/empty falls back to
+    the official ``hf-audio/xcodec-hubert-general-balanced``.
+    """
+    model = model or XCODEC1_DEFAULT_REPO
+    if os.path.isdir(model):
+        return model
+
+    from huggingface_hub import snapshot_download
+
+    try:
+        return snapshot_download(model)
+    except Exception as download_exc:
+        try:
+            return snapshot_download(model, local_files_only=True)
+        except Exception:
+            raise RuntimeError(
+                f"Could not resolve the XCodec1 checkpoint {model!r}: the download failed and "
+                "no cached snapshot was found. Download it manually or point the TTA deploy "
+                "yaml's stage-1 model at a local copy."
+            ) from download_exc
+
+
+def ensure_audiogen_weights(model_dir: str) -> None:
+    """Link index-referenced shards missing from ``model_dir`` from the full checkpoint.
+
+    No-op when the shards are already present (or the layout is unexpected);
+    raises only if a referenced shard exists nowhere.
+    """
+    index_path = os.path.join(model_dir, "model.safetensors.index.json")
+    if not os.path.isfile(index_path):
+        return
+    try:
+        with open(index_path) as f:
+            weight_map = json.load(f).get("weight_map", {})
+    except (OSError, ValueError) as exc:
+        logger.warning("Could not read %s: %s", index_path, exc)
+        return
+
+    source_dir = os.path.join(os.path.dirname(os.path.abspath(model_dir)), _WEIGHT_SOURCE_FOLDER)
+    for shard in sorted(set(weight_map.values())):
+        dst = os.path.join(model_dir, shard)
+        if os.path.exists(dst):
+            continue
+        src = os.path.join(source_dir, shard)
+        if not os.path.exists(src):
+            raise FileNotFoundError(
+                f"Audex checkpoint shard {shard!r} is missing from {model_dir} and "
+                f"{source_dir}. Download the repo's {_WEIGHT_SOURCE_FOLDER}/ folder "
+                "(it holds the deduplicated weight shards)."
+            )
+        try:
+            os.symlink(os.path.relpath(src, model_dir), dst)
+            logger.info("Linked Audex weight shard %s -> %s", dst, src)
+        except OSError:
+            shutil.copy2(src, dst)
+            logger.info("Copied Audex weight shard %s -> %s", src, dst)

vllm_omni/model_executor/models/audex/pipeline.py

@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Audex (Nemotron-Labs-Audex-2B) TTS pipeline topology.
+
+Stage 0: Thinker  — ChatML text prompt → <speechcodec_N> tokens (LLM AR).
+Stage 1: Code2Wav — streaming causal speech decoder → 16 kHz waveform.
+
+Users pass the HF repo ROOT (its manifest config.json carries
+``model_type: nemotron_labs_audex``); per-stage checkpoints resolve through
+``model_subdir``/``tokenizer_subdir``. The decoder subfolder has no tokenizer
+files, so stage 1's tokenizer also points at the thinker checkpoint.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.audex"
+
+# <speechgen_end> token id in checkpoint_folder_audiogen's tokenizer; pinned
+# by a unit test against the real tokenizer fixture.
+AUDEX_SPEECHGEN_END_TOKEN_ID = 131076
+
+# <audiogen_end> stops TTA generation (pinned in tests via models/audex/tta.py).
+AUDEX_AUDIOGEN_END_TOKEN_ID = 131074
+
+# No pipeline-level model_arch: stage-0 classes resolve from each
+# checkpoint's own ``architectures`` (dense on the 2B, NemotronH on the
+# 30B-A3B); forcing one name via hf_overrides would misbind the other
+# size. Decoder stages keep their explicit per-stage model_arch.
+AUDEX_TTS_PIPELINE = PipelineConfig(
+    model_type="audex_tts",
+    default_deploy_config_name="audex_tts.yaml",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="audex_thinker",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            model_subdir="checkpoint_folder_audiogen",
+            tokenizer_subdir="checkpoint_folder_audiogen",
+            # Emits the unconditional CFG companion for cfg_scale > 1.0
+            # requests; a no-op for unguided requests.
+            prompt_expand_func=f"{_PROC}.expand_cfg_prompts",
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.thinker2code2wav_async_chunk"),
+            custom_process_next_stage_input_func=f"{_PROC}.thinker2code2wav_full_payload",
+            sampling_constraints={
+                "detokenize": False,
+                "stop_token_ids": [AUDEX_SPEECHGEN_END_TOKEN_ID],
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="audex_code2wav",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="AudexCode2Wav",
+            model_subdir="audex_causal_speech_decoder",
+            tokenizer_subdir="checkpoint_folder_audiogen",
+            sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only",
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)
+
+
+# Text-to-audio (caption → general audio): same audiogen thinker checkpoint,
+# but the output rides the interleaved 4-codebook <audiocodec_N> RVQ block
+# and decodes through the EXTERNAL XCodec1 checkpoint. XCodec1 is a CNN codec
+# decoded over the full sequence, so this pipeline is sync full-payload only
+# (deploy yaml sets async_chunk: false). CFG is effectively mandatory for TTA
+# quality (official default scale 3.0).
+# Cascaded speech-to-speech: the audio-capable full checkpoint serves all
+# three official passes (ASR → chat → TTS) in one deployment; only the TTS
+# pass (output modality "audio", final_stage_id 1) traverses the streaming
+# speech decoder — text-modality passes finish at stage 0 and never reach
+# the codec path. The cascade itself is orchestrated by the client/example.
+AUDEX_S2S_PIPELINE = PipelineConfig(
+    model_type="audex_s2s",
+    default_deploy_config_name="audex_s2s.yaml",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="audex_omni",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="latent",
+            model_subdir="checkpoint_folder_full",
+            tokenizer_subdir="checkpoint_folder_full",
+            prompt_expand_func=f"{_PROC}.expand_cfg_prompts",
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.thinker2code2wav_async_chunk"),
+            custom_process_next_stage_input_func=f"{_PROC}.thinker2code2wav_full_payload",
+            # <speechgen_end> additionally stops the TTS pass; text passes
+            # stop at their natural eos first. detokenize stays on for the
+            # text-final passes.
+            sampling_constraints={
+                "detokenize": True,
+                "stop_token_ids": [AUDEX_SPEECHGEN_END_TOKEN_ID],
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="audex_code2wav",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="AudexCode2Wav",
+            model_subdir="audex_causal_speech_decoder",
+            # The decoder folder has no tokenizer files. Borrow the audiogen
+            # tokenizer (not checkpoint_folder_full): this stage resolves its
+            # snapshot with the "tts" download profile, which guarantees
+            # checkpoint_folder_audiogen but not the full folder's tokenizer,
+            # so per-stage / clean-cache launches would otherwise fail.
+            tokenizer_subdir="checkpoint_folder_audiogen",
+            sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only",
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)
+
+
+# Audio understanding (ASR / audio QA): single-stage deployment of the
+# audio-capable full checkpoint — speech (+ text instruction) in, text out.
+# Modeled on the ming_flash_omni / qwen2_5_omni thinker-only pipelines.
+AUDEX_THINKER_ONLY_PIPELINE = PipelineConfig(
+    model_type="audex_thinker_only",
+    default_deploy_config_name="audex_thinker_only.yaml",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="audex_omni",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="text",
+            model_subdir="checkpoint_folder_full",
+            tokenizer_subdir="checkpoint_folder_full",
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)
+
+
+AUDEX_TTA_PIPELINE = PipelineConfig(
+    model_type="audex_tta",
+    default_deploy_config_name="audex_tta.yaml",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="audex_tta_thinker",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            model_subdir="checkpoint_folder_audiogen",
+            tokenizer_subdir="checkpoint_folder_audiogen",
+            prompt_expand_func=f"{_PROC}.expand_cfg_prompts",
+            custom_process_next_stage_input_func=f"{_PROC}.thinker2xcodec_full_payload",
+            sampling_constraints={
+                "detokenize": False,
+                "stop_token_ids": [AUDEX_AUDIOGEN_END_TOKEN_ID],
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="audex_xcodec",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="AudexXCodec1",
+            # Length placeholder for the stage-1 request; the bulk codec
+            # payload ships via thinker2xcodec_full_payload.
+            sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only",
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)

vllm_omni/model_executor/models/audex/prompt.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Audex TTS/TTA prompt construction.
+
+The Audex thinker consumes the exact ChatML prompt used by the official
+inference script (``inference_scripts_vllm/audiogen_scripts/run_audio_gen_vllm.py``
+in the nvidia/Nemotron-Labs-Audex-2B repo). The checkpoint's bundled chat
+template opens a thinking block (``<think>\\n``) instead of the closed
+``<think></think>`` + ``<speechgen_start>`` priming that TTS generation
+requires, so the prompt is built from a literal template here; a unit test
+pins it byte-for-byte against the official format.
+
+``build_null_prompt`` is the unconditional-prompt counterpart used by
+classifier-free guidance: the transcription is replaced with repeated
+``<unk>`` tokens, iteratively adjusted so the tokenized null prompt is
+exactly as long as the conditional prompt (the CFG pair must decode the
+same positions).
+"""
+
+AUDEX_SYSTEM_PROMPT = "You are a helpful and harmless assistant.\n\nYou are not allowed to use any tools."
+
+# ``{text}`` carries its own leading whitespace: the conditional prompt uses
+# " {transcription}" (matching the official script byte-for-byte), while the
+# CFG null prompt inserts the bare ``<unk>`` run. ``<unk>`` is a special
+# token, so a leading space would tokenize as a standalone space token and
+# make some conditional lengths unreachable by whole-``<unk>`` steps.
+_TTS_PROMPT_TEMPLATE = (
+    "<|im_start|>system\n{system_prompt}<|im_end|>\n"
+    "<|im_start|>user\n<|text to speech|> Generate speech for this transcription.{text}<|im_end|>\n"
+    "<|im_start|>assistant\n<think></think><speechgen_start>"
+)
+
+_TTA_PROMPT_TEMPLATE = (
+    "<|im_start|>system\n{system_prompt}<|im_end|>\n"
+    "<|im_start|>user\n<|text to audio|> Generate audio for this caption.{text}<|im_end|>\n"
+    "<|im_start|>assistant\n<think></think><audiogen_start>"
+)
+
+
+def build_cond_prompt(text: str) -> str:
+    """Build the conditional TTS prompt for one transcription."""
+    text = text.strip()
+    if not text:
+        raise ValueError("Audex TTS requires non-empty input text")
+    return _TTS_PROMPT_TEMPLATE.format(system_prompt=AUDEX_SYSTEM_PROMPT, text=f" {text}")
+
+
+def build_tta_cond_prompt(caption: str) -> str:
+    """Build the conditional TTA prompt for one audio caption."""
+    caption = caption.strip()
+    if not caption:
+        raise ValueError("Audex TTA requires a non-empty caption")
+    return _TTA_PROMPT_TEMPLATE.format(system_prompt=AUDEX_SYSTEM_PROMPT, text=f" {caption}")
+
+
+def _length_matched_null(cond_prompt: str, tokenizer, template: str, max_iters: int) -> str:
+    target_len = len(tokenizer.encode(cond_prompt))
+    n_unk = 1
+    for _ in range(max_iters):
+        prompt = template.format(system_prompt=AUDEX_SYSTEM_PROMPT, text="<unk>" * n_unk)
+        diff = target_len - len(tokenizer.encode(prompt))
+        if diff == 0:
+            return prompt
+        n_unk = max(1, n_unk + diff)
+        if diff < 0 and n_unk == 1:
+            break
+    raise ValueError(f"Could not length-match the CFG null prompt to {target_len} tokens")
+
+
+def build_null_prompt(cond_prompt: str, tokenizer, max_iters: int = 64) -> str:
+    """Unconditional TTS prompt for CFG (length-matched ``<unk>`` padding).
+
+    Replaces the transcription with ``<unk>`` repeated, adjusting the count
+    until ``tokenizer.encode`` yields exactly the conditional prompt's
+    length. Raises when no count matches so callers can fail the request
+    instead of submitting a misaligned pair.
+    """
+    return _length_matched_null(cond_prompt, tokenizer, _TTS_PROMPT_TEMPLATE, max_iters)
+
+
+def build_tta_null_prompt(cond_prompt: str, tokenizer, max_iters: int = 64) -> str:
+    """Unconditional TTA prompt for CFG (length-matched ``<unk>`` padding)."""
+    return _length_matched_null(cond_prompt, tokenizer, _TTA_PROMPT_TEMPLATE, max_iters)

vllm_omni/model_executor/models/audex/speech_decoder/__init__.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+"""Vendored Audex causal speech decoder (Apache-2.0).
+
+Verbatim copy of the ``audex_causal_speech_decoder`` remote-code package from
+nvidia/Nemotron-Labs-Audex-2B. Vendored because transformers' dynamic module
+loader resolves HuggingFace-cache blob symlinks to their real paths, which
+breaks the package's relative imports when loading with trust_remote_code
+directly from a hub snapshot.
+"""
+
+from .configuration_audex_causal_speech_decoder import AudexCausalSpeechDecoderConfig
+from .modeling_audex_causal_speech_decoder import (
+    AudexCausalSpeechDecoderModel,
+    AudexCausalSpeechDecoderSession,
+)
+
+__all__ = [
+    "AudexCausalSpeechDecoderConfig",
+    "AudexCausalSpeechDecoderModel",
+    "AudexCausalSpeechDecoderSession",
+]

vllm_omni/model_executor/models/audex/speech_decoder/configuration_audex_causal_speech_decoder.py

@@ -0,0 +1,50 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class AudexCausalSpeechDecoderConfig(PretrainedConfig):
+    model_type = "audex_causal_speech_decoder"
+
+    def __init__(
+        self,
+        hidden_dim: int = 2048,
+        depth: int = 12,
+        heads: int = 32,
+        pos_meb_dim: int = 64,
+        hop_length: int = 320,
+        vq_dim: int = 2048,
+        lookahead_steps: int = 4,
+        sample_rate: int = 16000,
+        codebook_levels: list[int] | None = None,
+        codebook_size: int = 65536,
+        token_embed_dim: int = 8,
+        embed_tokens_from_codes: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = hidden_dim
+        self.depth = depth
+        self.heads = heads
+        self.pos_meb_dim = pos_meb_dim
+        self.hop_length = hop_length
+        self.vq_dim = vq_dim
+        self.lookahead_steps = lookahead_steps
+        self.sample_rate = sample_rate
+        self.codebook_levels = codebook_levels or [4, 4, 4, 4, 4, 4, 4, 4]
+        self.codebook_size = codebook_size
+        self.token_embed_dim = token_embed_dim
+        self.embed_tokens_from_codes = embed_tokens_from_codes

vllm_omni/model_executor/models/audex/speech_decoder/modeling_audex_causal_speech_decoder.py

@@ -0,0 +1,478 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from transformers import PreTrainedModel
+
+from .configuration_audex_causal_speech_decoder import AudexCausalSpeechDecoderConfig
+from .streaming_utils import load_audex_causal_speech_decoder as _load_decoder_for_remote_code
+
+REMOTE_CODE_IMPORTS = (_load_decoder_for_remote_code,)
+
+
+class RotaryPositionalEmbeddings(nn.Module):
+    def __init__(self, dim: int, max_seq_len: int = 4096, base: int = 10_000) -> None:
+        super().__init__()
+        self.dim = dim
+        self.base = base
+        self.max_seq_len = max_seq_len
+        self.rope_init()
+        self._rope_ready = False
+
+    def rope_init(self, device: torch.device | None = None) -> None:
+        theta = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device)[: (self.dim // 2)].float() / self.dim))
+        self.register_buffer("theta", theta, persistent=False)
+        self.build_rope_cache(self.max_seq_len)
+
+    def build_rope_cache(self, max_seq_len: int = 4096) -> None:
+        self.max_seq_len = max_seq_len
+        seq_idx = torch.arange(max_seq_len, dtype=self.theta.dtype, device=self.theta.device)
+        idx_theta = torch.einsum("i, j -> ij", seq_idx, self.theta).float()
+        cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
+        self.register_buffer("cache", cache, persistent=False)
+
+    def forward(self, x: torch.Tensor, *, input_pos: torch.Tensor | None = None) -> torch.Tensor:
+        seq_len = x.size(1)
+        needed_seq_len = seq_len if input_pos is None else int(input_pos.max().item()) + 1
+        if (
+            not getattr(self, "_rope_ready", False)
+            or self.theta.device != x.device
+            or needed_seq_len > self.cache.size(0)
+        ):
+            self.rope_init(device=x.device)
+            if needed_seq_len > self.cache.size(0):
+                self.build_rope_cache(max(needed_seq_len, self.cache.size(0) * 2))
+            self._rope_ready = True
+
+        rope_cache = self.cache[:seq_len] if input_pos is None else self.cache[input_pos]
+        xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+        rope_cache = rope_cache.view(-1, xshaped.size(1), 1, xshaped.size(3), 2)
+        x_out = torch.stack(
+            [
+                xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
+                xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
+            ],
+            -1,
+        )
+        return x_out.flatten(3).type_as(x)
+
+
+class CausalCodecDecoderCache:
+    def __init__(self) -> None:
+        self.key_values: dict[int, tuple[Tensor, Tensor]] = {}
+        self.position = 0
+
+    def input_positions(self, length: int, device: torch.device) -> Tensor:
+        return torch.arange(self.position, self.position + length, device=device).unsqueeze(0)
+
+    def update(self, layer_idx: int, key: Tensor, value: Tensor) -> tuple[Tensor, Tensor]:
+        if layer_idx in self.key_values:
+            prev_key, prev_value = self.key_values[layer_idx]
+            key = torch.cat([prev_key, key], dim=2)
+            value = torch.cat([prev_value, value], dim=2)
+        self.key_values[layer_idx] = (key, value)
+        return key, value
+
+    def advance(self, length: int) -> None:
+        self.position += length
+
+    def reset(self) -> None:
+        self.key_values.clear()
+        self.position = 0
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm_x = torch.mean(x**2, dim=-1, keepdim=True)
+        return x * torch.rsqrt(norm_x + self.eps) * self.weight
+
+
+class MLP(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(dim, 4 * dim, bias=False)
+        self.silu = nn.SiLU()
+        self.fc2 = nn.Linear(4 * dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fc2(self.silu(self.fc1(x)))
+
+
+class Attention(nn.Module):
+    def __init__(self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings, layer_idx: int):
+        super().__init__()
+        if dim % n_heads != 0:
+            raise ValueError(f"dim must be divisible by n_heads, got dim={dim}, n_heads={n_heads}")
+        self.n_heads = n_heads
+        self.layer_idx = layer_idx
+        self.rotary_embed = rotary_embed
+        self.c_attn = nn.Linear(dim, 3 * dim, bias=False)
+        self.c_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cache: CausalCodecDecoderCache | None = None,
+        input_pos: Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+        qkv = self.c_attn(x)
+        head_dim = qkv.size(-1) // (3 * self.n_heads)
+        qkv = qkv.view(batch_size, seq_len, 3, self.n_heads, head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+
+        q = self.rotary_embed(q.transpose(1, 2), input_pos=input_pos).transpose(1, 2)
+        k = self.rotary_embed(k.transpose(1, 2), input_pos=input_pos).transpose(1, 2)
+        if cache is None:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            if input_pos is None:
+                raise ValueError("input_pos is required when cache is set")
+            k, v = cache.update(self.layer_idx, k, v)
+            key_pos = torch.arange(k.size(2), device=x.device).view(1, 1, 1, -1)
+            attn_mask = key_pos <= input_pos.view(input_pos.size(0), 1, -1, 1)
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        return self.c_proj(y.transpose(1, 2).contiguous().view(batch_size, seq_len, -1))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings, layer_idx: int):
+        super().__init__()
+        self.att_norm = RMSNorm(dim)
+        self.ffn_norm = RMSNorm(dim)
+        self.att = Attention(dim=dim, n_heads=n_heads, rotary_embed=rotary_embed, layer_idx=layer_idx)
+        self.mlp = MLP(dim=dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cache: CausalCodecDecoderCache | None = None,
+        input_pos: Tensor | None = None,
+    ) -> torch.Tensor:
+        x = x + self.att(self.att_norm(x), cache=cache, input_pos=input_pos)
+        return x + self.mlp(self.ffn_norm(x))
+
+
+class PatchHead(nn.Module):
+    def __init__(self, dim: int, hop_length: int = 320):
+        super().__init__()
+        self.proj = nn.Linear(dim, hop_length, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.tanh(self.proj(x))
+        return x.reshape(x.size(0), 1, -1)
+
+
+class CausalVocosBackbone(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int = 2048,
+        depth: int = 12,
+        heads: int = 32,
+        pos_meb_dim: int = 64,
+    ):
+        super().__init__()
+        rotary_embed = RotaryPositionalEmbeddings(dim=pos_meb_dim)
+        self.transformers = nn.ModuleList(
+            [
+                TransformerBlock(dim=hidden_dim, n_heads=heads, rotary_embed=rotary_embed, layer_idx=idx)
+                for idx in range(depth)
+            ]
+        )
+        self.final_layer_norm = RMSNorm(hidden_dim)
+
+    def forward(self, x: torch.Tensor, cache: CausalCodecDecoderCache | None = None) -> torch.Tensor:
+        input_pos = None
+        if cache is not None:
+            input_pos = cache.input_positions(x.size(1), x.device).expand(x.size(0), -1)
+        for block in self.transformers:
+            x = block(x, cache=cache, input_pos=input_pos)
+        if cache is not None:
+            cache.advance(x.size(1))
+        return self.final_layer_norm(x)
+
+
+class CausalCodecDecoderVocos(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int = 2048,
+        depth: int = 12,
+        heads: int = 32,
+        pos_meb_dim: int = 64,
+        hop_length: int = 320,
+        vq_dim: int = 2048,
+        lookahead_steps: int = 0,
+    ):
+        super().__init__()
+        if lookahead_steps < 0:
+            raise ValueError(f"lookahead_steps must be >= 0, got {lookahead_steps}")
+
+        self.wav_proj = nn.Linear(hop_length, hidden_dim, bias=False)
+        self.fc_post_a = nn.Linear(vq_dim, hidden_dim, bias=False)
+        self.lookahead_steps = lookahead_steps
+        if lookahead_steps > 0:
+            self.lookahead_conv = nn.Conv1d(
+                hidden_dim,
+                hidden_dim,
+                kernel_size=lookahead_steps + 1,
+                padding=0,
+                groups=hidden_dim,
+                bias=False,
+            )
+            self.lookahead_act = nn.SiLU()
+            self.lookahead_proj = nn.Conv1d(hidden_dim, hidden_dim, kernel_size=1, bias=False)
+            nn.init.zeros_(self.lookahead_proj.weight)
+        else:
+            self.lookahead_conv = None
+            self.lookahead_act = None
+            self.lookahead_proj = None
+        self.backbone = CausalVocosBackbone(hidden_dim, depth, heads, pos_meb_dim)
+        self.head = PatchHead(hidden_dim, hop_length)
+
+    def _project_tokens(self, vq_emb: torch.Tensor) -> torch.Tensor:
+        return self.fc_post_a(vq_emb)
+
+    def _apply_lookahead(self, x: torch.Tensor) -> torch.Tensor:
+        if self.lookahead_conv is None:
+            return x
+        if self.lookahead_act is None or self.lookahead_proj is None:
+            raise RuntimeError("lookahead modules are not initialized")
+        h = F.pad(x.transpose(1, 2), (0, self.lookahead_steps))
+        h = self.lookahead_proj(self.lookahead_act(self.lookahead_conv(h)))
+        return x + h.transpose(1, 2)
+
+    def _apply_lookahead_window(self, x: torch.Tensor) -> torch.Tensor:
+        if self.lookahead_conv is None:
+            return x
+        if self.lookahead_act is None or self.lookahead_proj is None:
+            raise RuntimeError("lookahead modules are not initialized")
+        if x.size(1) <= self.lookahead_steps:
+            raise ValueError(f"lookahead window needs more than {self.lookahead_steps} frames, got {x.size(1)}")
+        h = self.lookahead_proj(self.lookahead_act(self.lookahead_conv(x.transpose(1, 2))))
+        return x[:, : h.size(2)] + h.transpose(1, 2)
+
+    def decode_cached(
+        self,
+        vq_emb: torch.Tensor,
+        cache: CausalCodecDecoderCache,
+        lookahead_vq_emb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        x = self._project_tokens(vq_emb)
+        if self.lookahead_steps > 0:
+            if lookahead_vq_emb is None:
+                lookahead_vq_emb = vq_emb.new_zeros(vq_emb.size(0), self.lookahead_steps, vq_emb.size(-1))
+            if lookahead_vq_emb.size(1) != self.lookahead_steps:
+                raise ValueError(
+                    f"lookahead_vq_emb must have {self.lookahead_steps} frames, got {lookahead_vq_emb.size(1)}"
+                )
+            lookahead_x = self._project_tokens(lookahead_vq_emb)
+            x = self._apply_lookahead_window(torch.cat([x, lookahead_x], dim=1))
+        x = self.backbone(x, cache=cache)
+        return self.head(x)
+
+    def forward(
+        self,
+        vq_emb: torch.Tensor,
+        patched_wav: torch.Tensor | None = None,
+        alpha: float = 0.0,
+    ) -> torch.Tensor:
+        x = self._project_tokens(vq_emb)
+        x = self._apply_lookahead(x)
+        if patched_wav is not None:
+            h = self.wav_proj(patched_wav)
+            mask = torch.bernoulli(
+                torch.full(
+                    (x.size(0), x.size(1), 1),
+                    min(max(alpha, 0.0), 1.0),
+                    device=x.device,
+                    dtype=x.dtype,
+                )
+            )
+            x = x + h * mask
+        return self.head(self.backbone(x))
+
+
+class AudexSpeechTokenEmbedder(nn.Module):
+    def __init__(
+        self,
+        output_dim: int,
+        token_embed_dim: int,
+        codebook_levels: Sequence[int],
+    ) -> None:
+        super().__init__()
+        if len(codebook_levels) != token_embed_dim:
+            raise ValueError(
+                f"token_embed_dim={token_embed_dim} must match codebook_levels length={len(codebook_levels)}"
+            )
+        self.codebook_levels = tuple(int(level) for level in codebook_levels)
+        self.project_out = nn.Linear(token_embed_dim, output_dim)
+
+    def forward(self, indices: torch.Tensor) -> torch.Tensor:
+        if indices.size(-1) != 1:
+            raise ValueError(f"indices last dimension must be 1, got {indices.size(-1)}")
+        levels = torch.tensor(self.codebook_levels, dtype=torch.long, device=indices.device)
+        basis = torch.cumprod(torch.cat([levels.new_ones(1), levels[:-1]]), dim=0)
+        level_indices = (indices.long() // basis) % levels
+        dtype = self.project_out.weight.dtype
+        codes = level_indices.to(dtype=dtype)
+        levels = levels.to(dtype=dtype)
+        codes = codes * (2.0 / (levels - 1.0)) - 1.0
+        return self.project_out(codes)
+
+    def get_output_from_indices(self, indices: torch.Tensor) -> torch.Tensor:
+        return self(indices)
+
+
+class AudexCausalSpeechDecoderModel(PreTrainedModel):
+    config_class = AudexCausalSpeechDecoderConfig
+    base_model_prefix = "module"
+    all_tied_weights_keys: dict[str, Any] = {}
+    Cache = CausalCodecDecoderCache
+
+    def __init__(self, config: AudexCausalSpeechDecoderConfig):
+        super().__init__(config)
+        self.audex_speech_token_embedder = AudexSpeechTokenEmbedder(
+            output_dim=config.vq_dim,
+            token_embed_dim=config.token_embed_dim,
+            codebook_levels=config.codebook_levels,
+        )
+        self.module = CausalCodecDecoderVocos(
+            hidden_dim=config.hidden_dim,
+            depth=config.depth,
+            heads=config.heads,
+            pos_meb_dim=config.pos_meb_dim,
+            hop_length=config.hop_length,
+            vq_dim=config.vq_dim,
+            lookahead_steps=config.lookahead_steps,
+        )
+
+    @property
+    def lookahead_steps(self) -> int:
+        return self.module.lookahead_steps
+
+    def create_cache(self) -> CausalCodecDecoderCache:
+        return CausalCodecDecoderCache()
+
+    def decode_cached(
+        self,
+        vq_emb: torch.Tensor,
+        cache: CausalCodecDecoderCache,
+        lookahead_vq_emb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.module.decode_cached(vq_emb, cache, lookahead_vq_emb=lookahead_vq_emb)
+
+    def create_session(
+        self,
+        *,
+        chunk_frames: int = 1,
+        sample_rate: int | None = None,
+        return_numpy: bool = True,
+    ) -> AudexCausalSpeechDecoderSession:
+        return AudexCausalSpeechDecoderSession(
+            decoder=self,
+            chunk_frames=chunk_frames,
+            sample_rate=sample_rate or self.config.sample_rate,
+            return_numpy=return_numpy,
+        )
+
+    def forward(
+        self,
+        vq_emb: torch.Tensor,
+        patched_wav: torch.Tensor | None = None,
+        alpha: float = 0.0,
+    ) -> torch.Tensor:
+        return self.module(vq_emb, patched_wav=patched_wav, alpha=alpha)
+
+
+class AudexCausalSpeechDecoderSession:
+    def __init__(
+        self,
+        decoder: AudexCausalSpeechDecoderModel,
+        *,
+        chunk_frames: int,
+        sample_rate: int,
+        return_numpy: bool,
+    ):
+        if chunk_frames <= 0:
+            raise ValueError(f"chunk_frames must be positive, got {chunk_frames}")
+        self.decoder = decoder
+        self.chunk_frames = chunk_frames
+        self.sample_rate = sample_rate
+        self.return_numpy = return_numpy
+        self.cache = decoder.create_cache()
+        self.buffer: list[list[int]] = []
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.decoder.parameters()).device
+
+    def reset(self) -> None:
+        self.cache = self.decoder.create_cache()
+        self.buffer.clear()
+
+    def push(self, token_frames: Sequence[Sequence[int]]) -> Iterator[tuple[int, Any]]:
+        self.buffer.extend(list(frame) for frame in token_frames)
+        yield from self._drain(flush=False)
+
+    def flush(self) -> Iterator[tuple[int, Any]]:
+        yield from self._drain(flush=True)
+
+    def _drain(self, *, flush: bool) -> Iterator[tuple[int, Any]]:
+        ready_frames = len(self.buffer) - self.decoder.lookahead_steps
+        while self.buffer and (flush or ready_frames >= self.chunk_frames):
+            emit_frames = min(self.chunk_frames, len(self.buffer)) if flush else self.chunk_frames
+            wav = self._decode_buffered_frames(emit_frames, flush=flush)
+            del self.buffer[:emit_frames]
+            ready_frames = len(self.buffer) - self.decoder.lookahead_steps
+            yield self.sample_rate, self._format_chunk(wav)
+
+    def _embed_speech_token_frames(self, token_frames: Sequence[Sequence[int]]) -> torch.Tensor:
+        indices = torch.tensor(token_frames, dtype=torch.long, device=self.device).unsqueeze(0)
+        return self.decoder.audex_speech_token_embedder.get_output_from_indices(indices)
+
+    def _decode_buffered_frames(self, emit_frames: int, *, flush: bool) -> torch.Tensor:
+        with torch.inference_mode():
+            vq_emb = self._embed_speech_token_frames(self.buffer[:emit_frames])
+            lookahead_vq_emb = None
+            if self.decoder.lookahead_steps > 0:
+                future_frames = self.buffer[emit_frames : emit_frames + self.decoder.lookahead_steps]
+                future_parts = []
+                if future_frames:
+                    future_parts.append(self._embed_speech_token_frames(future_frames))
+                missing_frames = self.decoder.lookahead_steps - len(future_frames) if flush else 0
+                if missing_frames > 0:
+                    future_parts.append(vq_emb.new_zeros(vq_emb.size(0), missing_frames, vq_emb.size(-1)))
+                lookahead_vq_emb = torch.cat(future_parts, dim=1) if future_parts else None
+            return self.decoder.decode_cached(vq_emb, self.cache, lookahead_vq_emb=lookahead_vq_emb)
+
+    def _format_chunk(self, wav: torch.Tensor) -> Any:
+        chunk = wav.squeeze().float().detach().cpu()
+        if not self.return_numpy:
+            return chunk
+
+        import numpy as np
+
+        return chunk.numpy().astype(np.float32, copy=False)

vllm_omni/model_executor/models/audex/speech_decoder/streaming_utils.py

@@ -0,0 +1,140 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator, Sequence
+from typing import TYPE_CHECKING, Any
+
+import torch
+from transformers import AutoModel
+
+if TYPE_CHECKING:
+    from .modeling_audex_causal_speech_decoder import AudexCausalSpeechDecoderModel
+
+
+def freeze_model(model: torch.nn.Module) -> torch.nn.Module:
+    model.eval()
+    for param in model.parameters():
+        param.requires_grad = False
+    return model
+
+
+def load_audex_causal_speech_decoder(
+    model_path: str,
+    *,
+    device: torch.device | str = "cuda",
+    dtype: torch.dtype | str | None = None,
+) -> AudexCausalSpeechDecoderModel:
+    kwargs = {"trust_remote_code": True}
+    if dtype is not None:
+        kwargs["torch_dtype"] = dtype
+    model = AutoModel.from_pretrained(model_path, **kwargs)
+    return freeze_model(model.to(device))
+
+
+def _decoder_device(decoder: AudexCausalSpeechDecoderModel) -> torch.device:
+    return next(decoder.parameters()).device
+
+
+def _embed_speech_token_frames(
+    decoder: AudexCausalSpeechDecoderModel,
+    token_frames: Sequence[Sequence[int]],
+    device: torch.device,
+) -> torch.Tensor:
+    indices = torch.tensor(token_frames, dtype=torch.long, device=device).unsqueeze(0)
+    return decoder.audex_speech_token_embedder.get_output_from_indices(indices)
+
+
+def decode_speech_token_frames(
+    token_frames: Iterable[Sequence[int]],
+    decoder: AudexCausalSpeechDecoderModel,
+    *,
+    device: torch.device | str | None = None,
+    chunk_frames: int = 1,
+) -> Iterator[torch.Tensor]:
+    if chunk_frames <= 0:
+        raise ValueError(f"chunk_frames must be positive, got {chunk_frames}")
+
+    device = torch.device(device) if device is not None else _decoder_device(decoder)
+    cache = decoder.create_cache()
+    buffer: list[list[int]] = []
+    lookahead_steps = decoder.lookahead_steps
+
+    for token_frame in token_frames:
+        buffer.append(list(token_frame))
+        while len(buffer) > lookahead_steps:
+            emit_frames = min(chunk_frames, len(buffer) - lookahead_steps)
+            wav = _decode_buffered_frames(decoder, cache, buffer, emit_frames, device, flush=False)
+            del buffer[:emit_frames]
+            yield wav
+
+    while buffer:
+        emit_frames = min(chunk_frames, len(buffer))
+        wav = _decode_buffered_frames(decoder, cache, buffer, emit_frames, device, flush=True)
+        del buffer[:emit_frames]
+        yield wav
+
+
+async def decode_speech_token_stream(
+    token_stream: AsyncIterable[Sequence[int]],
+    decoder: AudexCausalSpeechDecoderModel,
+    *,
+    device: torch.device | str | None = None,
+    chunk_frames: int = 1,
+) -> AsyncIterator[torch.Tensor]:
+    if chunk_frames <= 0:
+        raise ValueError(f"chunk_frames must be positive, got {chunk_frames}")
+
+    device = torch.device(device) if device is not None else _decoder_device(decoder)
+    cache = decoder.create_cache()
+    buffer: list[list[int]] = []
+    lookahead_steps = decoder.lookahead_steps
+
+    async for token_frame in token_stream:
+        buffer.append(list(token_frame))
+        while len(buffer) > lookahead_steps:
+            emit_frames = min(chunk_frames, len(buffer) - lookahead_steps)
+            wav = _decode_buffered_frames(decoder, cache, buffer, emit_frames, device, flush=False)
+            del buffer[:emit_frames]
+            yield wav
+
+    while buffer:
+        emit_frames = min(chunk_frames, len(buffer))
+        wav = _decode_buffered_frames(decoder, cache, buffer, emit_frames, device, flush=True)
+        del buffer[:emit_frames]
+        yield wav
+
+
+def _decode_buffered_frames(
+    decoder: AudexCausalSpeechDecoderModel,
+    cache: Any,
+    buffer: Sequence[Sequence[int]],
+    emit_frames: int,
+    device: torch.device,
+    *,
+    flush: bool,
+) -> torch.Tensor:
+    with torch.inference_mode():
+        vq_emb = _embed_speech_token_frames(decoder, buffer[:emit_frames], device)
+        lookahead_vq_emb = None
+        if decoder.lookahead_steps > 0:
+            future_frames = buffer[emit_frames : emit_frames + decoder.lookahead_steps]
+            future_parts = []
+            if future_frames:
+                future_parts.append(_embed_speech_token_frames(decoder, future_frames, device))
+            missing_frames = decoder.lookahead_steps - len(future_frames) if flush else 0
+            if missing_frames > 0:
+                future_parts.append(vq_emb.new_zeros(vq_emb.size(0), missing_frames, vq_emb.size(-1)))
+            lookahead_vq_emb = torch.cat(future_parts, dim=1) if future_parts else None
+        return decoder.decode_cached(vq_emb, cache, lookahead_vq_emb=lookahead_vq_emb)

vllm_omni/model_executor/models/audex/tta.py

@@ -0,0 +1,293 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0):
+#   inference_scripts_vllm/audiogen_scripts/tta_rvq_logits_processor.py
+#   inference_scripts_vllm/audiogen_scripts/run_audio_gen_vllm_rvq_logit_mask.py
+# Divergences from the official files are commented at the divergence site.
+"""Audex text-to-audio (TTA) token space: RVQ phase masking and helpers.
+
+TTA output is an interleaved 4-codebook RVQ stream over ``<audiocodec_N>``
+tokens: generated position ``p`` (counted from ``<audiogen_start>``) must
+come from codebook ``p % 4``, whose codec-id range is
+``[phase * 1024, (phase + 1) * 1024)``. The tokenizer defines 8192
+``<audiocodec_*>`` symbols; only codec ids 0..4095 (the 4 quantizers XCodec1
+decodes) are generated — the upper half stays disallowed.
+
+Unguided/unmasked sampling emits phase-invalid sequences that XCodec1 cannot
+decode, so :class:`TTARVQPhaseMaskLogitsProcessor` is required for validity,
+driven per request by ``SamplingParams.extra_args["tta_rvq"]``:
+
+    {
+        "phase_token_ids": [[...1024 ids...] x 4],
+        "start_tid": <audiogen_start id>,
+        "end_tid": <audiogen_end id>,
+        "codec_cap": <max codec tokens or None>,
+        "start_in_prompt": <bool>,
+    }
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import torch
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.v1.sample.logits_processor import (
+    BatchUpdate,
+    LogitsProcessor,
+    MoveDirectionality,
+)
+
+logger = init_logger(__name__)
+
+NEG_INF = float("-inf")
+XCODEC1_NUM_CODEBOOKS = 4
+XCODEC1_CODEBOOK_SIZE = 1024
+
+# Pinned by a unit test against the real checkpoint tokenizer.
+AUDEX_AUDIOGEN_START_TOKEN_ID = 131073
+AUDEX_AUDIOGEN_END_TOKEN_ID = 131074
+AUDEX_AUDIOCODEC_TOKEN_OFFSET = 196613
+AUDEX_AUDIOCODEC_VOCAB_SIZE = 8192
+
+
+def build_tta_phase_token_ids(tokenizer: Any) -> tuple[list[list[int]], int, int]:
+    """Group tokenizer ids of ``<audiocodec_N>`` into the 4 RVQ phases.
+
+    Phase ``p`` collects the tokenizer ids for codec ids
+    ``[p*1024, (p+1)*1024)``. Fails loudly on missing markers or an
+    incomplete codec vocab.
+
+    Returns ``(phase_token_ids, audiogen_start_tid, audiogen_end_tid)``.
+    """
+    vocab = tokenizer.get_vocab()
+    start_tid = vocab.get("<audiogen_start>")
+    end_tid = vocab.get("<audiogen_end>")
+    if start_tid is None or end_tid is None:
+        raise ValueError(f"Missing marker token ids: audiogen_start={start_tid}, audiogen_end={end_tid}")
+
+    phase_token_ids: list[list[int]] = []
+    for phase in range(XCODEC1_NUM_CODEBOOKS):
+        ids: list[int] = []
+        missing: list[int] = []
+        for codec_id in range(phase * XCODEC1_CODEBOOK_SIZE, (phase + 1) * XCODEC1_CODEBOOK_SIZE):
+            tok = vocab.get(f"<audiocodec_{codec_id}>")
+            if tok is None:
+                missing.append(codec_id)
+            else:
+                ids.append(tok)
+        if missing:
+            raise ValueError(
+                f"Incomplete codec vocab for RVQ phase {phase}: missing {len(missing)} codec ids (e.g. {missing[:5]})"
+            )
+        phase_token_ids.append(ids)
+
+    return phase_token_ids, int(start_tid), int(end_tid)
+
+
+def validate_rvq_phase(
+    codec_ids: list[int],
+    codebook_size: int = XCODEC1_CODEBOOK_SIZE,
+    num_codebooks: int = XCODEC1_NUM_CODEBOOKS,
+) -> dict[str, Any]:
+    """Check that ``codec_ids[i]`` falls in RVQ phase ``i % 4``'s codec range."""
+    mismatch_count = 0
+    first_mismatch = None
+    for i, codec_id in enumerate(codec_ids):
+        expected = i % num_codebooks
+        lo = expected * codebook_size
+        if not (lo <= codec_id < lo + codebook_size):
+            mismatch_count += 1
+            if first_mismatch is None:
+                first_mismatch = [i, int(codec_id), int(codec_id // codebook_size), int(expected)]
+    return {
+        "phase_valid": mismatch_count == 0,
+        "mismatch_count": mismatch_count,
+        "first_mismatch": first_mismatch,
+        "codec_count": len(codec_ids),
+    }
+
+
+class TTARVQPhaseMaskLogitsProcessor(LogitsProcessor):
+    """Mask logits to the RVQ-phase-valid token set for Audex TTA requests.
+
+    Per batch row:
+      ``_req[idx]``           -> {"start_tid", "end_tid", "codec_cap", "start_in_prompt"}
+      ``_output_tokens[idx]`` -> live reference to the request's output ids
+      ``_start_pos[idx]``     -> cached index of <audiogen_start> once seen;
+                                 -1 means the prompt already ended with it.
+
+    The disallow masks are vocab-sized boolean tensors built once on the
+    first ``apply`` (using the real ``logits.shape[-1]`` so vocab padding is
+    also forbidden).
+    """
+
+    def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool) -> None:
+        self.device = device
+        self._req: dict[int, dict[str, Any]] = {}
+        self._output_tokens: dict[int, list[int]] = {}
+        self._start_pos: dict[int, int] = {}
+
+        self._phase_token_ids: list[list[int]] | None = None
+        self._mask_end_tid: int | None = None
+
+        self._disallow_phase: torch.Tensor | None = None
+        self._disallow_phase0_end: torch.Tensor | None = None
+        self._disallow_except_end: torch.Tensor | None = None
+
+    def is_argmax_invariant(self) -> bool:
+        return False
+
+    def _reset(self) -> None:
+        self._req.clear()
+        self._output_tokens.clear()
+        self._start_pos.clear()
+
+    def update_state(self, batch_update: BatchUpdate | None) -> None:
+        if batch_update is None:
+            return
+
+        for idx in batch_update.removed:
+            self._req.pop(idx, None)
+            self._output_tokens.pop(idx, None)
+            self._start_pos.pop(idx, None)
+
+        if not self._req and batch_update.added:
+            self._reset()
+
+        for idx, params, _, output_token_ids in batch_update.added:
+            extra_args = params.extra_args if params else None
+            tta = extra_args.get("tta_rvq") if extra_args else None
+            if tta:
+                self._req[idx] = {
+                    "start_tid": tta["start_tid"],
+                    "end_tid": tta["end_tid"],
+                    "codec_cap": tta.get("codec_cap"),
+                    "start_in_prompt": bool(tta.get("start_in_prompt", False)),
+                }
+                self._output_tokens[idx] = output_token_ids
+                if self._req[idx]["start_in_prompt"]:
+                    self._start_pos[idx] = -1
+                else:
+                    self._start_pos.pop(idx, None)
+                if self._phase_token_ids is None:
+                    self._phase_token_ids = tta["phase_token_ids"]
+                    self._mask_end_tid = tta["end_tid"]
+            else:
+                self._req.pop(idx, None)
+                self._output_tokens.pop(idx, None)
+                self._start_pos.pop(idx, None)
+
+        if self._req:
+            for src, dst, direction in batch_update.moved:
+                src_req = self._req.pop(src, None)
+                dst_req = self._req.pop(dst, None)
+                src_tokens = self._output_tokens.pop(src, None)
+                dst_tokens = self._output_tokens.pop(dst, None)
+                src_pos = self._start_pos.pop(src, None)
+                dst_pos = self._start_pos.pop(dst, None)
+                if src_req is not None:
+                    self._req[dst] = src_req
+                if src_tokens is not None:
+                    self._output_tokens[dst] = src_tokens
+                if src_pos is not None:
+                    self._start_pos[dst] = src_pos
+                if direction == MoveDirectionality.SWAP:
+                    if dst_req is not None:
+                        self._req[src] = dst_req
+                    if dst_tokens is not None:
+                        self._output_tokens[src] = dst_tokens
+                    if dst_pos is not None:
+                        self._start_pos[src] = dst_pos
+
+    def _build_masks(self, vocab_size: int) -> None:
+        if self._phase_token_ids is None or self._mask_end_tid is None:
+            raise RuntimeError(
+                "TTARVQPhaseMaskLogitsProcessor: phase token ids missing; "
+                "extra_args['tta_rvq'] was not provided for a TTA request."
+            )
+        if len(self._phase_token_ids) != XCODEC1_NUM_CODEBOOKS:
+            raise RuntimeError(
+                f"TTARVQPhaseMaskLogitsProcessor: expected {XCODEC1_NUM_CODEBOOKS} phase id "
+                f"lists, got {len(self._phase_token_ids)}."
+            )
+
+        allow = torch.zeros((XCODEC1_NUM_CODEBOOKS, vocab_size), dtype=torch.bool, device=self.device)
+        for phase, ids in enumerate(self._phase_token_ids):
+            if not ids:
+                raise RuntimeError(f"TTARVQPhaseMaskLogitsProcessor: empty token id list for RVQ phase {phase}.")
+            idx = torch.as_tensor(ids, dtype=torch.long, device=self.device)
+            if int(idx.max()) >= vocab_size or int(idx.min()) < 0:
+                raise RuntimeError(
+                    f"TTARVQPhaseMaskLogitsProcessor: phase {phase} token id out of range for vocab_size={vocab_size}."
+                )
+            allow[phase].index_fill_(0, idx, True)
+
+        end_onehot = torch.zeros(vocab_size, dtype=torch.bool, device=self.device)
+        if self._mask_end_tid < 0 or self._mask_end_tid >= vocab_size:
+            raise RuntimeError(
+                f"TTARVQPhaseMaskLogitsProcessor: end token id out of range "
+                f"for vocab_size={vocab_size}: {self._mask_end_tid}."
+            )
+        end_onehot[self._mask_end_tid] = True
+
+        self._disallow_phase = ~allow
+        self._disallow_phase0_end = ~(allow[0] | end_onehot)
+        self._disallow_except_end = ~end_onehot
+        logger.info(
+            "TTARVQPhaseMaskLogitsProcessor: built masks vocab=%d end_tid=%d (pid=%d)",
+            vocab_size,
+            self._mask_end_tid,
+            os.getpid(),
+        )
+
+    def apply(self, logits: torch.Tensor) -> torch.Tensor:
+        if not self._req:
+            return logits
+
+        if self._disallow_phase is None:
+            self._build_masks(logits.shape[-1])
+
+        num_rows = logits.shape[0]
+        for idx, state in self._req.items():
+            # Rows beyond this step's logits (pure-prefill / partially
+            # scheduled steps) have nothing to mask; the persistent-batch
+            # index re-enters range on the decode steps that sample.
+            if idx >= num_rows:
+                continue
+            tokens = self._output_tokens.get(idx)
+            if tokens is None:
+                continue
+            n = len(tokens)
+
+            start_pos = self._start_pos.get(idx)
+            if start_pos is None:
+                start_tid = state["start_tid"]
+                for i in range(n):
+                    if tokens[i] == start_tid:
+                        start_pos = i
+                        self._start_pos[idx] = i
+                        break
+                if start_pos is None:
+                    continue
+
+            position = n - start_pos - 1
+            if position < 0:
+                continue
+
+            cap = state["codec_cap"]
+            if cap is not None and position >= cap:
+                disallow = self._disallow_except_end
+            else:
+                phase = position & (XCODEC1_NUM_CODEBOOKS - 1)
+                if phase == 0:
+                    disallow = self._disallow_phase0_end if position > 0 else self._disallow_phase[0]
+                else:
+                    disallow = self._disallow_phase[phase]
+
+            logits[idx].masked_fill_(disallow, NEG_INF)
+
+        return logits

vllm_omni/model_executor/models/registry.py

@@ -84,6 +84,31 @@ _OMNI_MODELS = {
         "cosyvoice3",
         "CosyVoice3Model",
     ),
+    "NemotronDenseForCausalLM": (
+        "audex",
+        "audex_thinker",
+        "NemotronDenseForCausalLM",
+    ),
+    "AudexCode2Wav": (
+        "audex",
+        "audex_code2wav",
+        "AudexCode2Wav",
+    ),
+    "AudexXCodec1": (
+        "audex",
+        "audex_xcodec",
+        "AudexXCodec1",
+    ),
+    "NemotronDenseAudexForConditionalGeneration": (
+        "audex",
+        "audex_omni",
+        "NemotronDenseAudexForConditionalGeneration",
+    ),
+    "NemotronHAudexForConditionalGeneration": (
+        "audex",
+        "audex_omni",
+        "NemotronHAudexForConditionalGeneration",
+    ),
     "OmniVoiceModel": (
         "omnivoice",
         "omnivoice",

vllm_omni/model_executor/stage_input_processors/audex.py

@@ -0,0 +1,356 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Audex thinker → code2wav stage input processors.
+
+The thinker is a plain token-autoregressive LM whose sampled stream contains
+``<speechcodec_N>`` tokens in a contiguous vocab block. These processors read
+the sampled token ids directly (no tokenizer, no multimodal outputs), map them
+to codec frame ids by subtracting the vocab offset, and ship single-codebook
+frame payloads to the code2wav stage.
+
+The downstream decoder is stateful-streaming (one session per request), so
+chunks carry ONLY new frames — no left-context replay — plus a ``finished``
+flag whose final (possibly empty) payload triggers the decoder's lookahead
+flush.
+"""
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
+from vllm_omni.inputs.data import OmniTokensPrompt
+from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt
+
+logger = init_logger(__name__)
+
+CFG_UNCOND_SUFFIX = "__cfg_uncond"
+
+# Vocab layout of nvidia/Nemotron-Labs-Audex-2B checkpoint_folder_audiogen:
+# <speechcodec_0> .. <speechcodec_65535> occupy a contiguous id block. The
+# deploy yaml can override both values; a unit test pins these defaults
+# against the real tokenizer.
+_DEFAULT_CODEC_TOKEN_OFFSET = 131077
+_DEFAULT_CODEC_VOCAB_SIZE = 65536
+
+# TTA <audiocodec_N> block (see models/audex/tta.py; pinned against the real
+# tokenizer there). Interleaved 4-codebook RVQ.
+_DEFAULT_AUDIOCODEC_TOKEN_OFFSET = 196613
+_DEFAULT_AUDIOCODEC_VOCAB_SIZE = 8192
+_TTA_NUM_CODEBOOKS = 4
+
+_STATE_KEY = "_audex_async_state"
+
+
+def _ensure_list(value: Any) -> list[int]:
+    if value is None:
+        return []
+    if isinstance(value, torch.Tensor):
+        return value.reshape(-1).tolist()
+    return list(value)
+
+
+def _connector_extra(transfer_manager: Any) -> dict[str, Any]:
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    if isinstance(raw_cfg, dict):
+        extra = raw_cfg.get("extra", raw_cfg)
+        return extra if isinstance(extra, dict) else {}
+    return {}
+
+
+def _codec_frames(tokens: list[int], codec_offset: int, codec_size: int) -> list[int]:
+    """Map ``<speechcodec_N>`` token ids to codec frame ids, dropping all other tokens."""
+    return [int(t) - codec_offset for t in tokens if codec_offset <= int(t) < codec_offset + codec_size]
+
+
+def _finished_meta(request_id: Any) -> dict[str, Any]:
+    """Terminal-payload meta shared by the sync full-payload producers."""
+    return {
+        "finished": torch.tensor(True, dtype=torch.bool),
+        "req_id": [request_id],
+    }
+
+
+def _finished_request_codes(
+    transfer_manager: Any,
+    request: Any,
+    offset_key: str,
+    size_key: str,
+    default_offset: int,
+    default_size: int,
+) -> tuple[list[int], Any]:
+    """Codec frame ids from a finished request's full token stream, plus its id."""
+    cfg = _connector_extra(transfer_manager)
+    codec_offset = int(cfg.get(offset_key, default_offset))
+    codec_size = int(cfg.get(size_key, default_size))
+    output_token_ids = _ensure_list(getattr(request, "output_token_ids", []))
+    codes = _codec_frames(output_token_ids, codec_offset, codec_size)
+    request_id = getattr(request, "external_req_id", None) or getattr(request, "request_id", None)
+    return codes, request_id
+
+
+def _empty_finished_payload(request_id: Any, codebooks: int) -> dict[str, Any]:
+    """Terminal payload for a failed request on the sync full-payload path.
+
+    Stage 1 decodes it to empty audio, which the serving layer's no-audio
+    guard turns into a request-level error. Raising here instead would be
+    swallowed by the payload hook (no terminal payload -> stage-1 stall).
+    """
+    return {
+        "codes": {"audio": torch.empty((0, codebooks), dtype=torch.long)},
+        "meta": _finished_meta(request_id),
+    }
+
+
+def expand_cfg_prompts(prompt: dict[str, Any] | str, sampling_params: Any) -> list[ExpandedPrompt]:
+    """Emit the unconditional CFG companion for a guided Audex TTS request.
+
+    A request opts into CFG by carrying ``extra_args`` with
+    ``cfg_scale > 1.0``, ``cfg_role="cond"``, ``cfg_pair_id`` and the
+    pre-built length-matched ``cfg_null_prompt`` (constructed where a
+    tokenizer is available: the serving layer or the offline example).
+    Unguided requests (no/1.0 scale) expand to nothing, keeping the
+    disabled path identical to the non-CFG flow.
+    """
+    extra_args = getattr(sampling_params, "extra_args", None) or {}
+    try:
+        cfg_scale = float(extra_args.get("cfg_scale", 1.0))
+    except (TypeError, ValueError):
+        return []
+    if cfg_scale <= 1.0 or extra_args.get("cfg_role") != "cond":
+        return []
+
+    pair_id = extra_args.get("cfg_pair_id")
+    null_prompt = extra_args.get("cfg_null_prompt")
+    if not pair_id or not isinstance(null_prompt, str) or not null_prompt:
+        raise ValueError(
+            "Audex CFG request is missing cfg_pair_id/cfg_null_prompt in extra_args; "
+            "the caller must build the length-matched null prompt before submission"
+        )
+
+    if isinstance(prompt, dict):
+        companion_prompt: dict[str, Any] | str = {**prompt, "prompt": null_prompt}
+        companion_prompt.pop("prompt_token_ids", None)
+    else:
+        companion_prompt = null_prompt
+
+    return [
+        ExpandedPrompt(
+            prompt=companion_prompt,
+            role="uncond",
+            request_id_suffix=CFG_UNCOND_SUFFIX,
+            sampling_params_override={
+                "extra_args": {
+                    "cfg_scale": cfg_scale,
+                    "cfg_role": "uncond",
+                    "cfg_pair_id": pair_id,
+                }
+            },
+        )
+    ]
+
+
+def _extract_new_codes(
+    request: Any,
+    state: dict[str, Any],
+    codec_offset: int,
+    codec_size: int,
+) -> list[int]:
+    """Return codec frame ids from tokens sampled since the last call."""
+    output_token_ids = _ensure_list(getattr(request, "output_token_ids", []))
+    new_tokens = output_token_ids[int(state.get("seen_len", 0)) :]
+    state["seen_len"] = len(output_token_ids)
+    return _codec_frames(new_tokens, codec_offset, codec_size)
+
+
+def thinker2code2wav_async_chunk(
+    transfer_manager: Any,
+    multimodal_output: dict[str, Any] | None,
+    request: Any,
+    is_finished: bool = False,
+) -> OmniPayloadStruct | None:
+    """Async-chunk producer: sampled token stream → per-chunk codec frames."""
+    del multimodal_output  # thinker emits none; codes come from the token stream
+    request_id = request.external_req_id
+    finished = bool(is_finished or request.is_finished())
+
+    cfg = _connector_extra(transfer_manager)
+    chunk_frames = int(cfg.get("codec_chunk_frames", 25))
+    initial_chunk_frames = int(cfg.get("initial_codec_chunk_frames", chunk_frames))
+    codec_offset = int(cfg.get("codec_token_offset", _DEFAULT_CODEC_TOKEN_OFFSET))
+    codec_size = int(cfg.get("codec_vocab_size", _DEFAULT_CODEC_VOCAB_SIZE))
+    if chunk_frames <= 0 or initial_chunk_frames <= 0:
+        raise ValueError(
+            f"Invalid codec chunk config: codec_chunk_frames={chunk_frames}, "
+            f"initial_codec_chunk_frames={initial_chunk_frames}"
+        )
+
+    request_state = transfer_manager.request_payload.get(request_id)
+    if not isinstance(request_state, dict) or _STATE_KEY not in request_state:
+        request_state = {
+            _STATE_KEY: {
+                "seen_len": 0,
+                "pending": [],
+                "emitted_chunks": 0,
+                "total_frames": 0,
+                "terminal_sent": False,
+            }
+        }
+        transfer_manager.request_payload[request_id] = request_state
+    state = request_state[_STATE_KEY]
+    if state.get("terminal_sent", False):
+        return None
+
+    pending: list[int] = state["pending"]
+    new_codes = _extract_new_codes(request, state, codec_offset, codec_size)
+    pending.extend(new_codes)
+    state["total_frames"] = int(state.get("total_frames", 0)) + len(new_codes)
+
+    if finished and state["total_frames"] == 0:
+        # Fail fast: the transfer adapter logs this and still delivers the
+        # terminal marker, and the serving layer turns the audio-less request
+        # into a request-level error instead of a silent empty stream.
+        state["terminal_sent"] = True
+        raise ValueError(f"Audex thinker produced no codec tokens for request {request_id}")
+
+    # Hold back at least one frame on non-terminal emissions so the terminal
+    # chunk always carries codes: the generation scheduler never runs a final
+    # forward for an empty terminal chunk, and the decoder's lookahead tail is
+    # only flushed inside a forward that sees the finished flag.
+    target = initial_chunk_frames if state["emitted_chunks"] == 0 else chunk_frames
+    if not finished and len(pending) <= target:
+        return None
+
+    if finished:
+        emit, state["pending"] = pending, []
+        state["terminal_sent"] = True
+    else:
+        emit, state["pending"] = pending[:target], pending[target:]
+
+    state["emitted_chunks"] += 1
+    finished_flag = torch.tensor(finished, dtype=torch.bool)
+    # Shape [frames, 1]: 2-D codec payloads ride the connector's payload
+    # channel into ``runtime_additional_information`` (1-D tensors are
+    # rerouted onto the token path and never reach the model's payload view).
+    return OmniPayloadStruct(
+        codes=CodesStruct(audio=torch.tensor(emit, dtype=torch.long).reshape(-1, 1)),
+        meta=MetaStruct(
+            finished=finished_flag,
+            # The chunk adapter strips ``meta.finished`` before it reaches the
+            # stage-1 model; ``stream_finished`` survives and is what
+            # AudexCode2Wav uses to trigger the lookahead flush.
+            stream_finished=finished_flag,
+            req_id=[request_id],
+        ),
+    )
+
+
+def thinker2code2wav_full_payload(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any] | None,
+    request: Any,
+) -> dict[str, Any] | None:
+    """Sync-path producer: ship the full codec sequence once at finish."""
+    del pooling_output  # thinker emits none; codes come from the token stream
+    codes, request_id = _finished_request_codes(
+        transfer_manager,
+        request,
+        "codec_token_offset",
+        "codec_vocab_size",
+        _DEFAULT_CODEC_TOKEN_OFFSET,
+        _DEFAULT_CODEC_VOCAB_SIZE,
+    )
+    if not codes:
+        # Do NOT raise: the sync payload hook swallows exceptions and returns
+        # None, so no terminal payload would reach stage 1 and the request
+        # would stall. An empty FINISHED payload lets stage 1 complete with
+        # empty audio and the serving-layer no-audio guard fail the request.
+        logger.error(
+            "Audex thinker produced no codec tokens for request %s; sending empty terminal payload", request_id
+        )
+        return _empty_finished_payload(request_id, codebooks=1)
+    return {"codes": {"audio": codes}, "meta": _finished_meta(request_id)}
+
+
+def thinker2xcodec_full_payload(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any] | None,
+    request: Any,
+) -> dict[str, Any] | None:
+    """TTA sync-path producer: ship the de-interleaved RVQ payload at finish.
+
+    Filters ``<audiocodec_N>`` ids from the sampled stream, truncates to whole
+    4-token frames, and ships a ``[frames, 4]`` codec-id payload (offsets NOT
+    subtracted; the XCodec1 stage removes the per-layer offset).
+    """
+    del pooling_output  # thinker emits none; codes come from the token stream
+    codes, request_id = _finished_request_codes(
+        transfer_manager,
+        request,
+        "audiocodec_token_offset",
+        "audiocodec_vocab_size",
+        _DEFAULT_AUDIOCODEC_TOKEN_OFFSET,
+        _DEFAULT_AUDIOCODEC_VOCAB_SIZE,
+    )
+    # Error paths do NOT raise: the sync payload hook swallows exceptions and
+    # returns None, so no terminal payload would reach stage 1 and the
+    # request would stall. An empty FINISHED payload lets stage 1 complete
+    # with empty audio and the serving-layer no-audio guard fail the request.
+    if not codes:
+        logger.error(
+            "Audex TTA thinker produced no codec tokens for request %s; sending empty terminal payload", request_id
+        )
+        return _empty_finished_payload(request_id, codebooks=_TTA_NUM_CODEBOOKS)
+    usable = len(codes) - (len(codes) % _TTA_NUM_CODEBOOKS)
+    if usable == 0:
+        logger.error(
+            "Audex TTA thinker produced fewer than one full RVQ frame for request %s; sending empty terminal payload",
+            request_id,
+        )
+        return _empty_finished_payload(request_id, codebooks=_TTA_NUM_CODEBOOKS)
+
+    # Validate the ACTUAL generated stream against the 4-codebook phase
+    # contract before decode: XCodec's per-layer offset removal clamps
+    # invalid ids into range, which would silently mask a phase violation.
+    from vllm_omni.model_executor.models.audex.tta import validate_rvq_phase
+
+    validity = validate_rvq_phase(codes[:usable])
+    if not validity["phase_valid"]:
+        logger.error(
+            "Audex TTA thinker produced a phase-invalid RVQ stream for request %s: "
+            "%d mismatched positions (first mismatch [index, codec_id, got_phase, expected_phase] = %s); "
+            "sending empty terminal payload",
+            request_id,
+            validity["mismatch_count"],
+            validity["first_mismatch"],
+        )
+        return _empty_finished_payload(request_id, codebooks=_TTA_NUM_CODEBOOKS)
+
+    frames = torch.tensor(codes[:usable], dtype=torch.long).reshape(-1, _TTA_NUM_CODEBOOKS)
+    return {"codes": {"audio": frames}, "meta": _finished_meta(request_id)}
+
+
+def thinker2code2wav_token_only(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+):
+    """Sync-side stage-1 engine-input builder (length placeholder).
+
+    The bulk codec payload ships via the worker connector from
+    ``thinker2code2wav_full_payload``; this only sizes the stage-1 request.
+    """
+    del prompt
+    engine_inputs: list[OmniTokensPrompt] = []
+    for source_output in source_outputs:
+        if not source_output.finished:
+            continue
+        output = source_output.outputs[0]
+        prompt_ids = _ensure_list(source_output.prompt_token_ids)
+        raw_output_ids = _ensure_list(output.cumulative_token_ids)
+        if raw_output_ids[: len(prompt_ids)] == prompt_ids:
+            raw_output_ids = raw_output_ids[len(prompt_ids) :]
+        engine_inputs.append(OmniTokensPrompt(prompt_token_ids=raw_output_ids))
+    return engine_inputs

vllm_omni/worker/gpu_ar_model_runner.py

@@ -321,13 +321,23 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             "Qwen3TTSTalkerForConditionalGeneration",
             "Qwen3TTSCode2Wav",
             "CosyVoice3Model",
+            "NemotronDenseForCausalLM",
+            "NemotronHForCausalLM",
             "DyninOmniForConditionalGeneration",
             "IndexTTS2TalkerForConditionalGeneration",
         }
-        if (
-            getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS
-            or get_stage_connector_role(self.model_config) is not None
-        ):
+        # The stage-level ``model_arch`` override may be blank so the class
+        # resolves from the checkpoint's own ``architectures`` (e.g. the Audex
+        # TTA/TTS thinker stages, which bind dense on the 2B and NemotronH on
+        # the 30B-A3B). Gate on the RESOLVED architectures, not only the raw
+        # override — otherwise a blank-override producer stage never creates
+        # its worker connector and the sync full-payload flush silently never
+        # runs (downstream stage starves on connector input).
+        stage_archs = set(getattr(self.model_config, "architectures", None) or ())
+        model_arch_override = getattr(self.model_config, "model_arch", None)
+        if model_arch_override:
+            stage_archs.add(model_arch_override)
+        if stage_archs & _OMNI_CONNECTOR_INIT_ARCHS or get_stage_connector_role(self.model_config) is not None:
             self.init_omni_connectors(
                 model_config=self.model_config,
                 kv_transfer_manager=self.kv_transfer_manager,

vllm_omni/worker/gpu_generation_model_runner.py

@@ -64,6 +64,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             "MiMoAudioModel",
             "Qwen3TTSTalkerForConditionalGeneration",
             "Qwen3TTSCode2Wav",
+            "AudexCode2Wav",
+            "AudexXCodec1",
             "CosyVoice3Model",
             "DyninOmniForConditionalGeneration",
             "IndexTTS2S2MelDecoder",