docs: update add-tts-model skill and contributing guide with single-stage patterns (#2806)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
This commit is contained in:
1 parent
77e8ba0a7b
commit
e8d22d4fa8
6 files changed
+901
-57
No files matched your search
@@ -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": <new_chunk_only>, ...}`
|
||||
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[<key>]` 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/<model_name>/`
|
||||
@@ -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_<model>.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 <changed-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_<model>.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 <changed>` 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)
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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/<model_name>/*.py \
|
||||
vllm_omni/entrypoints/openai/serving_speech.py \
|
||||
vllm_omni/model_executor/models/registry.py \
|
||||
tests/e2e/offline_inference/test_<model_name>.py \
|
||||
tests/e2e/online_serving/test_<model_name>.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 <your@email.com>`. Use
|
||||
`-s`:
|
||||
|
||||
```bash
|
||||
git commit -s -m "feat(<model>): add <Model> 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 <branch> --force-with-lease
|
||||
```
|
||||
@@ -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_<model>.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`.
|
||||
@@ -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_<your_model>.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_<model>_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
|
||||
|
||||
|
||||
Reference in new issue
Block a user