diff --git a/.claude/skills/add-tts-model/SKILL.md b/.claude/skills/add-tts-model/SKILL.md index e64e7e763..963ffb4f6 100644 --- a/.claude/skills/add-tts-model/SKILL.md +++ b/.claude/skills/add-tts-model/SKILL.md @@ -8,10 +8,90 @@ description: "Integrate a new text-to-speech model into vLLM-Omni from HuggingFa ## Overview ``` -HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph - (Phase 1) (Phase 2) (Phase 3) (Phase 4) (Phase 5) +HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph -> Pre-commit/DCO + (Phase 1) (Phase 2) (Phase 3) (Phase 4) (Phase 5) (Phase 6) ``` +Three architecture patterns are supported: + +- **Two-stage pipeline** (e.g. Qwen3-TTS, Fish Speech, CosyVoice3): AR + code-predictor → audio decoder, connected via async_chunk for low-latency + streaming. Use this for maximum performance. +- **Single-stage AR via generator** (e.g. MOSS-TTS-Nano): entire model runs + inside one AR worker, streaming audio chunks through a per-request + `inference_stream()` generator. Use this when the upstream model bundles AR + + codec inseparably. See [references/single-stage-ar.md](references/single-stage-ar.md). +- **Single-stage, vLLM-native base LM + side computation** (e.g. VoxCPM2): + the base language model runs under vLLM's PagedAttention as a normal AR + model; diffusion / VAE / side computations run outside vLLM and are + attached via the runner post-processing hook. This is a distinct pattern + from the generator approach above — do not confuse the two. + +The single-stage variants skip Phase 4 (async_chunk) but Phase 5 (CUDA graph) +is still encouraged for the inner AR loop. + +## Cross-Cutting Invariants + +These rules apply to every TTS model regardless of architecture (AR vs AR+diffusion, single-stage vs two-stage, codec-based vs VAE-based). They surface repeatedly across PRs — check them at the end of every phase. + +### I1. Streaming output contract + +Pick exactly one per-step semantics for `forward()` and document it in the docstring: + +- **Delta**: yield only new audio samples produced this step. Preferred — linear cost, low memory. +- **Cumulative**: re-decode from step 0 every call. O(N²); only acceptable if the codec has no streaming decode path. + +If you choose **delta**, verify the full emit→consolidate→consume chain: + +1. `forward()` returns `{"model_outputs": , ...}` +2. `_consolidate_multimodal_tensors()` in `vllm_omni/engine/output_processor.py` concatenates the audio key into one tensor at finish. If it skips the key (`continue`), offline consumers receive only the final chunk. See `output_processor.py` for the concrete list of handled modality keys. +3. Streaming consumers (SSE, Gradio) receive per-step deltas; offline consumers (`engine.generate()`) receive a single concatenated tensor. + +Cumulative-vs-delta mismatch is the most common silent bug — offline RTF benchmarks pass, but users hear replays or truncation. + +### I2. Multimodal output consumer hygiene + +`outputs[0].outputs[0].multimodal_output[]` can be any of `Tensor`, `list[Tensor]` (pre-consolidation snapshot), `np.ndarray`, or scalar. When writing tests, examples, and benchmarks: + +- **Never** use `dict.get("a") or dict.get("b")` on tensor values — Python evaluates the tensor's boolean, raising `RuntimeError: Boolean value of Tensor with more than one value is ambiguous`. Use explicit `if x is None` chains. +- Always defensively handle the list form: `if isinstance(x, list): x = torch.cat([t.reshape(-1) for t in x], dim=0)`. +- Assert `shape` / `dtype` / `duration` explicitly; do not rely on truthiness for presence checks. + +### I3. Hot-loop GPU discipline + +Inside any per-step model loop (AR decode, diffusion solver, CFM Euler, vocoder block loop): + +- No `tensor.item()`, `.cpu()`, or `.tolist()` — each triggers a GPU→CPU sync; at 10 steps × 60 frames × 4 ops that is 2400 syncs per request. +- Prefer `dst.copy_(src)` over `dst.fill_(src.item())` when writing a scalar tensor into a buffer. +- Prefer `torch.compile(Model.forward, fullgraph=False)` on the whole forward over per-submodule compile — fewer dispatch boundaries, larger fusion regions. Measure before choosing granularity. +- No Python-side control flow that depends on tensor values; use `torch.where` / masking instead. + +Profile first, optimize second. See the profiling docs / project memory for the trace-analysis workflow. + +### I4. Validation pyramid + +Offline RTF alone is necessary but not sufficient. Every new TTS model must pass all three: + +| Layer | Catches | Tool | +|-------|---------|------| +| Offline RTF / duration check | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e | +| Browser streaming playback | Delta/cumulative bugs, chunk boundary glitches, TTFP regressions | Gradio demo over `/v1/audio/speech?stream=true` | +| Concurrent requests | Per-request state leaks, codec window round-robin gaps | `max_num_seqs>1` smoke test with 4+ parallel prompts | + +Declaring a model "done" without all three has shipped regressions more than once. + +### I5. Per-request state is owned by the request, not the model + +If the model caches *anything* across `forward()` calls (streaming generators, codec buffers, sliding-window pads, CUDA graph state), key it by request ID: + +```python +self._state: dict[str, YourState] = {} # request_key → state +# fetch: request_key = str(info.get("_omni_req_id", "0")) +# free on finish: del self._state[request_key] +``` + +A shared buffer silently corrupts audio across concurrent requests — the symptom is crosstalk or truncation only under load. + ## Phase 1: HuggingFace Reference **Goal**: Understand the reference implementation and verify it produces correct audio. @@ -85,6 +165,48 @@ When audio output is wrong, check in this order: 6. **Tensor layout**: Codebook-major vs frame-major ordering 7. **dtype**: Float32 for codec decoders (autocast can corrupt audio) +### Streaming Correctness Rules (single-stage and two-stage) + +These bugs appear in almost every new TTS PR. Check all before the first push. See also the cross-cutting invariants I1 (output contract) and I5 (per-request state) above — the rules below are the Phase 2-specific instances of those invariants: + +- **Accumulate codes across AR steps** — each `forward()` appends new codes; do not reset between steps or audio will be truncated (fish speech: `fix: accumulate audio_codes across steps`) +- **Emit delta audio, not full waveform** — in streaming mode yield only the new chunk per step, not the re-decoded full waveform from step 0 (fish speech: `fix: emit delta audio not full waveform`) +- **All return paths must emit `model_outputs`** — if any early-return branch skips setting `model_outputs`, the serving layer silently drops that step's audio (fish speech: `fix: ensure ALL return paths emit model_outputs`) +- **Per-request state isolation** — for batched concurrent requests, key all state by request ID; a shared buffer corrupts audio across requests (fish speech: `fix: per-request vocode + delta emission`) +- **Codec tensor device** — move codec codes to the codec decoder's device before calling decode; mismatches cause silent CPU fallback or crashes (fish speech: `fix: use model device for CUDA stream`) +- **AR stage `max_num_seqs`** — set to at least 4 in production configs; for single-stage models this is the only stage. For two-stage models, Stage 0 (AR) needs `max_num_seqs ≥ 4` to pipeline concurrent requests; Stage 1 (codec decoder) typically uses `max_num_seqs: 1` intentionally. Default of 1 everywhere causes audio gaps under concurrency because the codec window round-robins across requests (RFC #2568) + +### Optional Dependency Handling + +Patch optional dependencies (`torchaudio` / `torchcodec` / `soundfile`) at +the top of `load_weights()`, not at module import. Failures to do so cause +cryptic errors only on environments missing the optional package — after +the model is already deployed. See +[references/optional-deps.md](references/optional-deps.md) for the full +pattern, signature constraints, and MOSS-TTS-Nano reference. + +### Single-Stage AR Pattern (alternative to two-stage) + +When the upstream model cannot be cleanly split into an AR stage and a +separate decoder, run the full pipeline inside a single AR worker and +stream audio through a per-request `inference_stream()` generator keyed by +`_omni_req_id`. Stage config must set `worker_type: ar`, +`engine_output_type: audio`, `final_output: true`, `is_comprehension: true`, +and `async_chunk: false` at the top level. Only extract params from +`additional_information` that you actually forward, or pre-commit fails +`ruff F841`. + +Full walkthrough with the complete `forward()` / `_create_stream_gen()` +skeleton and stage-config fields: +[references/single-stage-ar.md](references/single-stage-ar.md). For an +in-tree reference, look for any single-stage AR model under +`vllm_omni/model_executor/models/` — e.g. the MOSS-TTS-Nano integration when +it lands. + +**VoxCPM2 is a different pattern** and should not reuse this skeleton — it +runs the base LM under vLLM PagedAttention with external side-computation. +See `plan/voxcpm2_native_ar_design.md`. + ### Deliverables - Model files in `vllm_omni/model_executor/models//` @@ -98,10 +220,66 @@ When audio output is wrong, check in this order: ### Steps -1. **Register in `serving_speech.py`**: - - Add model stage name to `_TTS_MODEL_STAGES` set - - Add model detection flag (e.g., `_is_fish_speech`) - - Implement prompt builder method (e.g., `_build_fish_speech_prompt()`) +1. **Register in `serving_speech.py`** — add all 5 points in a **single commit**; + partial integration causes hard-to-debug failures. This file is modified by every + model PR and is the most common source of rebase conflicts — see conflict note below. + + **Point 1** — stage constant (near the top, alongside the other `_*_TTS_MODEL_STAGES` sets): + ```python + _YOUR_MODEL_TTS_MODEL_STAGES = {"your_stage_key"} + ``` + + **Point 2** — union into `_TTS_MODEL_STAGES`: + ```python + _TTS_MODEL_STAGES: set[str] = ( + ... + | _YOUR_MODEL_TTS_MODEL_STAGES + ) + ``` + + **Point 3** — model type detection in `_detect_tts_model_type()`: + ```python + if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES: + return "your_model" + ``` + + **Point 4** — validation dispatch in `_validate_tts_request()`: + ```python + if self._tts_model_type == "your_model": + return self._validate_your_model_request(request) + ``` + + **Point 5** — validation + parameter-builder methods: + ```python + def _validate_your_model_request(self, request) -> str | None: + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + return None + + def _build_your_model_params(self, request) -> dict: + params = {"text": [request.input]} + if request.voice is not None: + params["voice"] = [request.voice] + return params + ``` + Wire `_build_your_model_params` into `_create_tts_request()` alongside the other + model-specific param builders. + + > **Two dispatch patterns coexist**: Fish Speech uses a `self._is_fish_speech` boolean + > instance attribute checked before `elif self._is_tts`, while all newer models + > (CosyVoice3, MOSS-TTS-Nano) use the `_tts_model_type` string returned by + > `_detect_tts_model_type()`. For new models, always use the `_tts_model_type` string + > pattern — do not add new `_is_*` flags. + + > **Unused variable rule**: only extract fields in `_build_your_model_params` that + > are actually forwarded to the model. Unused extractions fail `ruff F841`. + > For voice-cloning fields (`ref_audio` → `prompt_audio_path`, `ref_text` → + > `prompt_text`), add them to the param builder and verify they reach the model call. + + **Rebase conflict note**: when rebasing onto `main` after another model was merged, + `serving_speech.py` will conflict. Resolution: always keep *both* the upstream + model's additions and your own — never discard either side. + 2. **Handle model-specific parameters**: - Voice cloning: `ref_audio` encoding and prompt injection - `max_new_tokens` override in sampling params @@ -129,11 +307,29 @@ def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list: ### Deliverables -- Updated `serving_speech.py` with model-specific prompt builder +- Updated `serving_speech.py` with all 5 integration points (single commit) - Client scripts and server launcher - Gradio demo with streaming and voice cloning UI +- E2E online serving test (`tests/e2e/online_serving/test_.py`) +- Buildkite CI entry in `.buildkite/test-merge.yml` - Documentation (offline + online serving docs) +### E2E test pitfalls to avoid + +- **One `OmniServerParams` set per file.** `omni_server` is module-scoped; a second + id in the same file forces mid-module teardown/restart and exposes startup + races (`APIConnectionError` on the first request post-restart). Split variants + into separate files instead. +- **No external URL fetches from the server.** CI and some dev hosts can't + reach `raw.githubusercontent.com` over TLS. Inline ref audio as + `data:audio/wav;base64,...`; the serving layer accepts both URL and data URL. +- **Use the harness readiness gate.** The fixture waits for HTTP 200 on + `/health`; don't add `time.sleep` in tests. If warmup is incomplete, make + `/health` return non-200 until you're actually ready. +- **Mark with `@pytest.mark.core_model` + `hardware_test(res={"cuda": "H100"})`** + so the test lands in `test-ready.yml` (triggered by the `ready` label) rather + than only nightly. + ## Phase 4: Async Chunk (Streaming) **Goal**: Enable inter-stage streaming so audio chunks are produced while AR generation continues. @@ -199,35 +395,10 @@ Stage 0 (AR) Stage 1 (Decoder) - Fixed batch size (fall back to eager for other sizes) - No dynamic control flow inside the graph -### Example: Code Predictor CUDA Graph (Qwen3-TTS) - -```python -import torch - -class CodePredictorGraph: - """Captures the 16-step code predictor AR loop as a single CUDA graph.""" - - def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64): - self.num_steps = 16 - self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device) - self.positions = torch.arange(self.num_steps, device=device) - self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device)) - self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device) - self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long) - # Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...) - - def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor: - self.input_buf.copy_(initial_input) - self.graph.replay() - return self.output_buf.clone() -``` - -### Performance Expectations - -Based on Qwen3-TTS code predictor experience: -- **3-5x speedup** for the graphed component -- Only effective for fixed batch sizes (typically batch_size=1) -- Falls back to eager mode for unsupported configurations +See [references/cuda-graph-example.md](references/cuda-graph-example.md) for +a worked skeleton (Qwen3-TTS code predictor, 16-step AR loop), performance +expectations (3–5× on the graphed component for fixed batch_size=1), and the +graph-safety constraints you must honor inside the captured region. ### Deliverables @@ -235,10 +406,33 @@ Based on Qwen3-TTS code predictor experience: - Benchmark script comparing eager vs graph performance - Documentation of constraints and fallback behavior +## Phase 6: Pre-commit and DCO + +**Goal**: Every commit passes `pre-commit` lint and carries a DCO +`Signed-off-by` line that matches the author email. + +- Install hooks once: `pre-commit install`. +- Run `pre-commit run --files ` before every push; accept any + auto-fixes, stage, re-commit. +- Sign every commit with `git commit -s`. DCO checks that author email and + `Signed-off-by` email match — `git config user.email` must match your + GitHub account email. + +Common pre-commit failures, recovery commands for missing sign-off, and the +full `pre-commit run` invocation for a TTS model: +[references/precommit-dco.md](references/precommit-dco.md). + ## Integration Checklist Use this checklist when integrating a new TTS model: +### Cross-Cutting Invariants (verify at end of every phase) +- [ ] I1: `forward()` docstring states cumulative vs delta; consolidation path audited end-to-end +- [ ] I2: Tests / examples / benchmarks never use `dict.get(a) or dict.get(b)` on tensor values; list form handled +- [ ] I3: No `.item()` / `.cpu()` / Python branch on tensor values inside per-step loops +- [ ] I4: Offline RTF, browser streaming playback, and concurrent-request smoke test all pass +- [ ] I5: Any cross-step cache keyed by `_omni_req_id`; entries freed when the request finishes + ### Phase 1: HF Reference - [ ] Reference model runs and produces correct audio - [ ] Architecture documented (stages, codebooks, tokens, sample rate) @@ -248,17 +442,27 @@ Use this checklist when integrating a new TTS model: - [ ] Model registered in `registry.py` - [ ] Config classes created with `model_type` registration - [ ] Stage 0 (AR) implemented and generates correct tokens -- [ ] Stage 1 (Decoder) produces correct audio from tokens +- [ ] Stage 1 (Decoder) produces correct audio from tokens — dtype float32 for codec decoder +- [ ] Stage 1 `max_num_seqs` ≥ 4 in production config (default 1 causes gaps under concurrency) +- [ ] Optional dependency fallbacks handled at `load_weights()` time (torchaudio/soundfile/etc.) +- [ ] Streaming: codec codes accumulated across AR steps (not reset per step) +- [ ] Streaming: delta audio emitted per chunk, not full re-decoded waveform +- [ ] Streaming: all `forward()` return paths emit `model_outputs` +- [ ] Streaming: per-request state keyed by request ID (not shared across requests) +- [ ] Streaming: codec tensors moved to codec decoder device before decode - [ ] Stage config YAML created - [ ] `end2end.py` produces audio matching reference quality - [ ] README.md written ### Phase 3: Online Serving -- [ ] Model added to `serving_speech.py` +- [ ] All 5 `serving_speech.py` integration points added in one commit +- [ ] Only extract params in `_build_*_params` that are forwarded to the model call (ruff F841) - [ ] Prompt builder handles text input correctly - [ ] Voice cloning works (if supported) - [ ] All response formats work (wav, mp3, flac, pcm) - [ ] Client scripts and server launcher created +- [ ] E2E online serving test written (`tests/e2e/online_serving/test_.py`) +- [ ] Buildkite CI entry added to `.buildkite/test-merge.yml` - [ ] Gradio demo working - [ ] Documentation added (offline + online docs, nav, supported models) @@ -276,9 +480,25 @@ Use this checklist when integrating a new TTS model: - [ ] Benchmark shows meaningful speedup - [ ] Fallback to eager works for unsupported configs +### Phase 6: Pre-commit and DCO +- [ ] `pre-commit run --files ` passes before every push +- [ ] Every commit has `Signed-off-by` matching the author email (`git commit -s`) +- [ ] `git config user.email` matches the email registered on your GitHub account +- [ ] Details and failure-recovery commands: [references/precommit-dco.md](references/precommit-dco.md) + ## References -- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) -- supported models and usage -- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) -- complete example of Phases 1-3 -- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) -- complete example of all 5 phases +In-skill references (details split out of the main body): + +- [references/single-stage-ar.md](references/single-stage-ar.md) — full `forward()` / generator skeleton for the MOSS-TTS-Nano-style pattern +- [references/optional-deps.md](references/optional-deps.md) — torchaudio / torchcodec fallback pattern +- [references/cuda-graph-example.md](references/cuda-graph-example.md) — Qwen3-TTS code-predictor CUDA graph skeleton +- [references/precommit-dco.md](references/precommit-dco.md) — full pre-commit invocation, failure table, DCO recovery + +Project docs and adjacent skills: + +- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) — supported models and usage +- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) — complete example of Phases 1–3 +- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) — complete example of all 5 phases - [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md) +- `plan/voxcpm2_native_ar_design.md` — VoxCPM2's vLLM-native AR + side-computation pattern (distinct from the generator-based single-stage described above) diff --git a/.claude/skills/add-tts-model/references/cuda-graph-example.md b/.claude/skills/add-tts-model/references/cuda-graph-example.md new file mode 100644 index 000000000..6f4993b5c --- /dev/null +++ b/.claude/skills/add-tts-model/references/cuda-graph-example.md @@ -0,0 +1,42 @@ +# CUDA Graph Example: Qwen3-TTS Code Predictor + +Reference sketch for capturing the 16-step code-predictor AR loop as a single +CUDA graph. Adapt the shapes, number of steps, and KV-head layout to your +model. + +```python +import torch + +class CodePredictorGraph: + """Captures the 16-step code predictor AR loop as a single CUDA graph.""" + + def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64): + self.num_steps = 16 + self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device) + self.positions = torch.arange(self.num_steps, device=device) + self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device)) + self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device) + self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long) + # Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...) + + def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor: + self.input_buf.copy_(initial_input) + self.graph.replay() + return self.output_buf.clone() +``` + +## Performance expectations (Qwen3-TTS code predictor) + +- **3–5× speedup** on the graphed component. +- Effective only for fixed batch sizes (typically `batch_size=1`). +- Fall back to eager for any shape/config that wasn't captured — do not try + to recapture per request. + +## Graph-safety constraints + +- `torch.argmax` instead of `torch.multinomial`. +- Fixed batch size. +- No Python control flow that branches on tensor values inside the captured + region (use `torch.where` / masks). +- No `.item()`, `.cpu()`, `.tolist()` — each would break the capture or + cause a GPU→CPU sync during replay. diff --git a/.claude/skills/add-tts-model/references/optional-deps.md b/.claude/skills/add-tts-model/references/optional-deps.md new file mode 100644 index 000000000..0a55f30f0 --- /dev/null +++ b/.claude/skills/add-tts-model/references/optional-deps.md @@ -0,0 +1,47 @@ +# Optional Dependency Handling + +Models that rely on `torchaudio`, `torchcodec`, `soundfile`, or other optional +packages must handle the missing-package case at import time, not at call +time. Failing to do this causes cryptic errors only on environments without +the optional package — after the model is already deployed. + +## Pattern (used in MOSS-TTS-Nano) + +```python +def _patch_torchaudio_load() -> None: + """Fallback torchaudio.load/save to soundfile if torchcodec is unavailable.""" + try: + import torchcodec # noqa: F401 + return # torchcodec present, torchaudio works as-is + except ImportError: + pass + + import soundfile as sf + import torchaudio + + def _sf_load(path, **kwargs): + data, sr = sf.read(str(path), dtype="float32", always_2d=True) + return torch.from_numpy(data).T, sr + + torchaudio.load = _sf_load + # patch .save similarly if needed +``` + +## Rules + +- Mirror the full signature of the replaced function. `torchaudio.load` + accepts `frame_offset`, `num_frames`, `normalize`, `channels_first`, + `format` — missing any of them causes `TypeError` from calling code. +- Catch `except Exception`, not just `ImportError`. `import torchaudio` + itself can fail with non-`ImportError` errors on broken installs. Log the + exception type and message (`logger.warning("torchaudio probe failed: %s: %s", + type(exc).__name__, exc)`) before falling back, so unrelated errors are not + silently swallowed. +- Call the patch function at the top of `load_weights()` before loading any + audio assets. Do not call it at module import time. + +## Reference implementation + +Any in-tree model that patches `torchaudio.load` in its `load_weights()` — +e.g. MOSS-TTS-Nano's `modeling_moss_tts_nano.py` once that integration +lands. diff --git a/.claude/skills/add-tts-model/references/precommit-dco.md b/.claude/skills/add-tts-model/references/precommit-dco.md new file mode 100644 index 000000000..86a1f42ce --- /dev/null +++ b/.claude/skills/add-tts-model/references/precommit-dco.md @@ -0,0 +1,54 @@ +# Pre-commit and DCO + +Every commit must pass `pre-commit` lint and carry a `Signed-off-by` line +that matches the commit author email. + +## Pre-commit + +Install hooks once: + +```bash +pre-commit install +``` + +Run before every push on the files you changed: + +```bash +pre-commit run --files \ + vllm_omni/model_executor/models//*.py \ + vllm_omni/entrypoints/openai/serving_speech.py \ + vllm_omni/model_executor/models/registry.py \ + tests/e2e/offline_inference/test_.py \ + tests/e2e/online_serving/test_.py +``` + +When pre-commit **modifies files** (ruff format auto-fix), it exits non-zero +but the changes are correct — stage the modified files and re-commit. + +| Failure | Root cause | Fix | +|---------|-----------|-----| +| `ruff F841` | Variable extracted but never forwarded to model call | Remove the extraction or wire it through | +| `ruff E402` | Import added below function definitions | Move to top-level import block | +| `ruff format` | Line length, spacing, quote style | Accept auto-fix, stage, re-commit | + +## DCO sign-off + +Every commit must carry `Signed-off-by: Your Name `. Use +`-s`: + +```bash +git commit -s -m "feat(): add TTS support" +``` + +Or set it permanently: `git config format.signOff true`. + +The DCO check verifies that the commit author email matches the +`Signed-off-by` line. Confirm `git config user.email` matches your GitHub +account email before committing. + +Fix a missing or mismatched sign-off on the latest commit: + +```bash +git commit --amend -s --no-edit +git push origin --force-with-lease +``` diff --git a/.claude/skills/add-tts-model/references/single-stage-ar.md b/.claude/skills/add-tts-model/references/single-stage-ar.md new file mode 100644 index 000000000..ed53d3026 --- /dev/null +++ b/.claude/skills/add-tts-model/references/single-stage-ar.md @@ -0,0 +1,108 @@ +# Single-Stage AR Pattern + +When the upstream model cannot be cleanly split into an AR stage and a separate +decoder (e.g. MOSS-TTS-Nano, or any model that bundles AR + codec via an +`inference_stream()` generator), run the whole pipeline inside a single AR +worker that yields audio chunks per request. + +This is distinct from VoxCPM2's pattern, which also runs in a single stage but +uses vLLM's native PagedAttention on the base language model with diffusion / +VAE side-computation outside vLLM — see +`plan/voxcpm2_native_ar_design.md` for that variant. + +## Implementation + +1. **Single model file** — load both AR LM and codec inside + `modeling_.py`. +2. **Load weights in `load_weights()`**, not `__init__()` — vLLM initializes + distributed state before any CUDA allocations. +3. **Stream via a per-request generator** stored in `self._stream_gens`: + +```python +class YourModelForCausalLM(nn.Module): + def __init__(self, *, vllm_config, prefix=""): + super().__init__() + self._lm = None # populated in load_weights() + self._stream_gens: dict = {} # request_key → generator + + def load_weights(self, weights): + # Load self._lm here, after vLLM distributed init + ... + + def forward( + self, + input_ids, + positions, + intermediate_tensors=None, + inputs_embeds=None, + runtime_additional_information: list[dict] | None = None, # one dict per request + **kwargs, + ) -> OmniOutput: + infos = runtime_additional_information or [{}] + # Skip dummy/profiling calls + if not runtime_additional_information or all(i.get("_is_dummy") for i in infos): + self._ar_emit_stop_token = True + return OmniOutput(...) # return empty outputs + + outputs, last_flags = [], [] + for info in infos: + request_key = str(info.get("_omni_req_id", "0")) # per-request ID from vLLM + if request_key not in self._stream_gens: + self._stream_gens[request_key] = self._create_stream_gen(info) + try: + chunk, is_last = next(self._stream_gens[request_key]) + except StopIteration: + chunk, is_last = torch.zeros(0), True + if is_last: + del self._stream_gens[request_key] + outputs.append(chunk) + last_flags.append(is_last) + + self._ar_emit_stop_token = all(last_flags) + return OmniOutput(multimodal_outputs={"model_outputs": outputs, ...}) + + def _create_stream_gen(self, info: dict): + """Yield (waveform_tensor, is_last) tuples from inference_stream().""" + for event in self._lm.inference_stream(...): + if event["type"] == "audio": + yield event["waveform"], False + elif event["type"] == "result": + # Fallback: some models emit a single "result" event instead of + # incremental "audio" events — handle both paths + yield event.get("waveform", torch.zeros(0)), True + return + yield torch.zeros(0), True + + def compute_logits(self, hidden_states, sampling_metadata): + # Emit EOS only after the last chunk so the AR scheduler ends the request + ... +``` + +## Key points + +- `runtime_additional_information` is the correct parameter name (not + `**kwargs`) — it carries one dict per request in the batch. +- The request ID is `info.get("_omni_req_id")` — set by vLLM, not by user code. +- Handle both `"audio"` (incremental) and `"result"` (final combined) event + types from upstream models. + +## Stage config + +Single stage with `worker_type: ar`, `engine_output_type: audio`, +`final_output: true`, `is_comprehension: true`, and `async_chunk: false` at +the top level. Omitting any of these causes silent misclassification in the +serving layer. + +## Lint discipline + +Only extract variables from `additional_information` that you actually +forward to the model call — unused extractions trip `ruff F841` in +pre-commit. + +## Reference implementation + +Look for any single-stage AR model under +`vllm_omni/model_executor/models/` — e.g. `moss_tts_nano/` when its +integration lands. If none is in tree yet, follow the skeleton above and +cross-check against the `is_comprehension: true` / `async_chunk: false` +dispatch in `vllm_omni/entrypoints/openai/serving_speech.py`. diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 622064173..34fd2dbb5 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -1,20 +1,93 @@ # Adding a TTS Model -This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS** -as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the -key optimizations all TTS models in this repo should follow. +This guide walks through adding a new TTS model to vLLM-Omni. Two patterns are +supported: + +- **Two-stage pipeline** (e.g. Qwen3-TTS, Fish Speech): an AR code-predictor stage + feeds an audio decoder stage via the `async_chunk` framework. This is the standard + pattern for maximum streaming performance. +- **Single-stage AR model** (e.g. MOSS-TTS-Nano): the model runs entirely inside one + AR worker and streams audio chunks directly from its own `inference_stream()` generator. + +Qwen3-TTS is used as the reference for the two-stage pattern. For the single-stage +pattern, refer to MOSS-TTS-Nano. ## Table of Contents 1. [Overview](#overview) -2. [Directory Structure](#directory-structure) -3. [Step-by-Step Implementation](#step-by-step-implementation) -4. [Key Components](#key-components) -5. [Model Registration](#model-registration) -6. [Stage Configuration](#stage-configuration) -7. [Stage Input Processors](#stage-input-processors) -8. [Testing](#testing) -9. [Summary](#summary) +2. [Cross-Cutting Invariants](#cross-cutting-invariants) +3. [Directory Structure](#directory-structure) +4. [Step-by-Step Implementation](#step-by-step-implementation) +5. [Key Components](#key-components) +6. [Model Registration](#model-registration) +7. [Stage Configuration](#stage-configuration) +8. [Stage Input Processors](#stage-input-processors) +9. [Online Serving Integration](#online-serving-integration) +10. [Single-Stage Models](#single-stage-models) +11. [Testing](#testing) +12. [Pre-commit and DCO](#pre-commit-and-dco) +13. [Summary](#summary) + +## Cross-Cutting Invariants + +These rules apply to every TTS model regardless of architecture (AR vs AR+diffusion, +single-stage vs two-stage, codec-based vs VAE-based). Each has surfaced as a silent +bug in a shipped PR — check them at the end of every phase, not just at the start. + +**I1. Streaming output contract.** Pick one per-step semantics for `forward()` and +document it in the docstring: + +- *Delta*: yield only new audio samples produced this step. Preferred — linear cost. +- *Cumulative*: re-decode from step 0 every call. O(N²); only acceptable when the + codec exposes no streaming decode. + +If you choose delta, audit the full chain: `forward()` returns the new chunk → +`_consolidate_multimodal_tensors()` in `vllm_omni/engine/output_processor.py` +concatenates the audio key into a single tensor at finish → streaming consumers +receive per-step chunks, offline consumers receive the concatenated tensor. A +mismatch (consolidator skips the key with `continue`, or consumers expect a list +but receive a tensor) is invisible in offline RTF benchmarks — users hear replays +or truncation only under live playback. + +**I2. Multimodal output consumer hygiene.** `outputs[0].outputs[0].multimodal_output[key]` +can be `Tensor`, `list[Tensor]` (pre-consolidation snapshot), `np.ndarray`, or +scalar. In every test, example, and benchmark: + +- Never write `dict.get("a") or dict.get("b")` on tensor values — Python evaluates + the tensor's truthiness and raises `Boolean value of Tensor with more than one + value is ambiguous`. Use explicit `if x is None` chains. +- Defensively handle the list form: + `if isinstance(x, list): x = torch.cat([t.reshape(-1) for t in x], dim=0)`. +- Assert `shape` / `dtype` / `duration` explicitly — do not rely on truthiness for + presence checks. + +**I3. Hot-loop GPU discipline.** Inside any per-step model loop (AR decode, +diffusion solver, CFM Euler step, per-frame vocoder): + +- No `tensor.item()`, `.cpu()`, or `.tolist()` — each triggers a GPU→CPU sync; a + 10-step × 60-frame × 4-op loop creates 2400 syncs per request. +- Prefer `dst.copy_(src)` over `dst.fill_(src.item())` for scalar-into-buffer writes. +- Whole-model `torch.compile(Model.forward, fullgraph=False)` usually outperforms + per-submodule compile — fewer dispatch boundaries, larger fusion regions. Measure + before choosing granularity. +- No Python control flow that depends on tensor values; use `torch.where` or masking. + +Profile before optimizing. + +**I4. Validation pyramid.** Offline RTF alone is necessary but not sufficient. A +new TTS model must pass all three levels: + +| Layer | Catches | Tool | +|-------|---------|------| +| Offline RTF / duration | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e | +| Browser streaming playback | Delta-vs-cumulative bugs, chunk boundary glitches, TTFP regressions | Gradio demo over `/v1/audio/speech?stream=true` | +| Concurrent requests | Per-request state leaks, codec window round-robin gaps | `max_num_seqs>1` smoke with 4+ parallel prompts | + +**I5. Per-request state belongs to the request.** If the model caches anything +across `forward()` calls (streaming generators, codec buffers, sliding-window pads, +CUDA graph state), key it by `info.get("_omni_req_id")` and free the entry on +request finish. A shared buffer silently corrupts audio across concurrent requests — +the symptom is crosstalk or truncation under load, nothing in single-request tests. ## Overview @@ -556,6 +629,302 @@ Recommended test cases for a new TTS model: Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py` +### E2E Online Serving Tests (`tests/e2e/online_serving/test_.py`) + +The `omni_server` fixture in `tests/conftest.py` is **module-scoped**. Each distinct +`OmniServerParams` id in the same test file forces the fixture to tear the server +down and spawn a new one mid-module. A few rules that save real CI debugging time: + +- **Prefer a single `OmniServerParams` set per file.** If you need to exercise two + deploy variants (e.g. `model.yaml` and `model_async_chunk.yaml`), either use one + variant and exercise streaming via request args, or split into two test files so + each file does exactly one server lifecycle. Mid-module teardown/restart is the + fragile path and surfaces startup races first. +- **Never depend on server-side fetching of external URLs** for reference audio or + other fixture data. CI runners (and China-hosted dev boxes) routinely fail on + SSL/DNS for public URLs. Inline the payload as a `data:audio/wav;base64,...` + ref_audio value — the serving layer accepts both forms. +- **Don't roll your own readiness probe.** The harness already waits for HTTP 200 + on `/health` before releasing the server to the test. If your model needs extra + warmup signals, expose them through `/health` rather than adding `time.sleep(...)` + inside the test. (Bare TCP `connect_ex` probes were insufficient; see + `tests/conftest.py` `OmniServer.wait_for_ready`.) +- **Use `core_model` marker + H100 hardware_test** to match the `test-ready.yml` + pipeline so your test is picked up by the `ready` label, not only nightly. + +## Online Serving Integration + +To expose your model through the `/v1/audio/speech` OpenAI-compatible endpoint, add +**all five** of the following integration points to +`vllm_omni/entrypoints/openai/serving_speech.py` in a **single commit**. Adding them +piecemeal causes partial-integration failures that are hard to debug. + +### 1. Stage constant + +Near the top of the file, alongside the other `_*_TTS_MODEL_STAGES` constants: + +```python +_YOUR_MODEL_TTS_MODEL_STAGES = {"your_model_stage_key"} +``` + +### 2. Union into `_TTS_MODEL_STAGES` + +Add to the `_TTS_MODEL_STAGES` set union: + +```python +_TTS_MODEL_STAGES: set[str] = ( + ... + | _YOUR_MODEL_TTS_MODEL_STAGES +) +``` + +### 3. Model type detection + +In `_detect_tts_model_type()`, add before the final `return None`: + +```python +if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES: + return "your_model" +``` + +### 4. Request validation dispatch + +In `_validate_tts_request()`, add before the fallback `return`: + +```python +if self._tts_model_type == "your_model": + return self._validate_your_model_request(request) +``` + +### 5. Validation and parameter-builder methods + +Add two new methods: + +```python +def _validate_your_model_request( + self, request: OpenAICreateSpeechRequest +) -> str | None: + """Validate YourModel request. Returns an error string or None.""" + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + return None + +def _build_your_model_params( + self, request: OpenAICreateSpeechRequest +) -> dict[str, Any]: + """Build additional_information dict for YourModel.""" + params: dict[str, Any] = {"text": [request.input]} + if request.voice is not None: + params["voice"] = [request.voice] + # Add any other model-specific fields here + return params +``` + +Then wire `_build_your_model_params` into the request-dispatch block in +`_create_tts_request()` (search for the equivalent `_build_*_params` call for an +existing model to find the right location). If the model supports voice cloning +(`ref_audio` → `prompt_audio_path`, `ref_text` → `prompt_text`), add those mappings +here too — follow any existing `_build__params` in `serving_speech.py` (e.g. +`_build_moss_tts_params` for the voice-cloning variant) for the pattern. + +> **Two dispatch patterns coexist:** Fish Speech uses a `self._is_fish_speech` boolean +> checked *before* `elif self._is_tts`. All newer models use the `_tts_model_type` +> string pattern shown above. For new models, always use the string pattern — do not +> add new `_is_*` boolean flags. + +> **Note on unused variables:** Only extract parameters in `_build_your_model_params` +> that you actually pass to the model's generate / `inference_stream` call. Extracting +> a variable without forwarding it will trigger a `ruff F841` pre-commit failure. + +### Merge conflicts + +`serving_speech.py` is modified by every new model PR and is the most common source of +rebase conflicts. When rebasing onto `main` and a conflict appears here, the resolution +is always to **keep both** the upstream model's additions and your own — never discard +either side. After resolving: + +```bash +git add vllm_omni/entrypoints/openai/serving_speech.py +git rebase --continue +``` + +## Single-Stage Models + +Some TTS models (e.g. MOSS-TTS-Nano) do not use a two-stage pipeline. Instead the +entire AR LM and audio decoder run inside a single AR worker, streaming audio chunks +directly from the model's own generator. + +### Directory structure + +``` +vllm_omni/model_executor/models/your_model_name/ + __init__.py + modeling_your_model_name.py # unified class: load_weights + forward + streaming + +vllm_omni/model_executor/stage_configs/your_model_name.yaml +``` + +No stage input processor is needed. + +### Stage config + +Use a single stage with `worker_type: ar`. The `is_comprehension: true` field and the +top-level `async_chunk: false` are required — omitting them causes silent +misclassification in the serving layer. Set `max_num_seqs` to at least 4 for +concurrent production use. + +```yaml +# stage_configs/your_model_name.yaml +async_chunk: false + +stage_args: + - stage_id: 0 + stage_type: llm + is_comprehension: true # required for serving_speech.py dispatch + runtime: + devices: "0" + engine_args: + model_stage: your_model_stage_key + model_arch: YourModelForCausalLM + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + engine_output_type: audio + max_num_seqs: 4 # min 4 for concurrent requests; default 1 causes gaps + final_output: true + final_output_type: audio +``` + +### Generator-based streaming pattern + +This is the MOSS-TTS-Nano pattern, distinct from VoxCPM2's vLLM-native AR pattern +(see `plan/voxcpm2_native_ar_design.md` for that variant). Load model weights in +`load_weights()` (not `__init__`) so vLLM finishes distributed initialisation before +any CUDA allocations. Stream via a per-request generator stored in an instance dict: + +```python +class YourModelForCausalLM(nn.Module): + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self._lm: nn.Module | None = None # populated in load_weights() + self._stream_gens: dict[str, Any] = {} # request_key → generator + + def load_weights(self, weights): + # Load self._lm here, after vLLM distributed init + ... + + def forward( + self, + input_ids, + positions, + intermediate_tensors=None, + inputs_embeds=None, + runtime_additional_information: list[dict] | None = None, # one dict per request + **kwargs, + ) -> OmniOutput: + infos = runtime_additional_information or [{}] + # Return empty output during dummy/profiling calls + if not runtime_additional_information or all(i.get("_is_dummy") for i in infos): + self._ar_emit_stop_token = True + return OmniOutput(...) + + outputs, last_flags = [], [] + for info in infos: + request_key = str(info.get("_omni_req_id", "0")) # set by vLLM, not user code + if request_key not in self._stream_gens: + self._stream_gens[request_key] = self._create_stream_gen(info) + try: + chunk, is_last = next(self._stream_gens[request_key]) + except StopIteration: + chunk, is_last = torch.zeros(0), True + if is_last: + del self._stream_gens[request_key] + outputs.append(chunk) + last_flags.append(is_last) + + self._ar_emit_stop_token = all(last_flags) + return OmniOutput(multimodal_outputs={"model_outputs": outputs, "is_last": last_flags}) + + def _create_stream_gen(self, info: dict): + """Yield (waveform_tensor, is_last) from the model's inference_stream(). + + Handle both incremental ("audio" events) and batch ("result" event) models: + some upstream implementations emit one "result" event with the full waveform + instead of incremental "audio" events. Both paths must be covered. + """ + for event in self._lm.inference_stream(...): + if event["type"] == "audio": + yield event["waveform"], False + elif event["type"] == "result": + # Fallback for models that don't emit incremental audio events + yield event.get("waveform", torch.zeros(0)), True + return + yield torch.zeros(0), True + + def compute_logits(self, hidden_states, sampling_metadata): + # Emit EOS only when the last chunk has been yielded so the AR + # scheduler ends the request at the right time. + ... +``` + +For an in-tree reference, look for any single-stage AR model under +`vllm_omni/model_executor/models/` (for example +`moss_tts_nano/modeling_moss_tts_nano.py` once its integration has landed). + +## Pre-commit and DCO + +All contributions must pass the pre-commit checks and the Developer Certificate of +Origin (DCO) sign-off before merging. + +### Running pre-commit + +Install the hooks once with `pre-commit install`. Then run before committing: + +```bash +pre-commit run --files \ + vllm_omni/model_executor/models/your_model_name/*.py \ + vllm_omni/entrypoints/openai/serving_speech.py \ + vllm_omni/model_executor/models/registry.py \ + tests/e2e/offline_inference/test_your_model_name.py \ + tests/e2e/online_serving/test_your_model_name.py +``` + +When pre-commit **modifies files**, it exits with a non-zero code but the reformatting +is correct. Stage the modified files and commit again — do not revert the changes. + +Common failures and fixes: + +| Check | Cause | Fix | +|-------|-------|-----| +| `ruff F841` | Local variable assigned but never used | Remove the extraction or forward it to the model call | +| `ruff E402` | Module-level import not at top of file | Move import to the top-level import block | +| `ruff format` | Line length, spacing, or quote style | Accept the auto-fix, stage, and re-commit | + +### DCO sign-off + +Every commit must carry a `Signed-off-by` trailer. Use the `-s` flag when committing: + +```bash +git commit -s -m "feat(your-model): add YourModel TTS support" +``` + +Or configure git to add it automatically: + +```bash +git config format.signOff true +``` + +To fix a missing sign-off on the most recent commit: + +```bash +git commit --amend -s --no-edit +git push origin your-branch --force-with-lease +``` + +> The DCO check verifies that the commit author email matches the `Signed-off-by` email. +> Make sure `git config user.email` is set to the address associated with your GitHub +> account before committing. + ## Adding a Model Recipe After implementing and testing your model, add a model recipe to the @@ -567,15 +936,19 @@ for the expected format. Adding a TTS model to vLLM-Omni involves: -1. **Create model directory** with AR stage, decoder stage, and unified class +1. **Create model directory** with AR stage, decoder stage, and unified class (two-stage) + or a single unified class with generator-based streaming (single-stage) 2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly 3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()` 4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration` 5. **Register** all stage classes in `registry.py` -6. **YAML configs** - provide both batch and `async_chunk` variants -7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25 -8. **Tests** - cover single request, batching, and async_chunk streaming -9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes) +6. **YAML configs** - provide both batch and `async_chunk` variants (two-stage), or a single-stage AR config +7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25 (two-stage only) +8. **Online serving** - add all 5 integration points to `serving_speech.py` in one commit +9. **Tests** - cover single request, batching, and streaming +10. **Pre-commit + DCO** - run `pre-commit` before pushing; sign every commit with `git commit -s` +11. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes) +12. **Invariants** - re-check I1–I5 (streaming contract, consumer hygiene, hot-loop discipline, validation pyramid, per-request state) at the end of every phase ### Qwen3-TTS Reference Files