Add EasyMagpie vLLM-Omni serving

examples/tts/easymagpie_vllm_omni/PERFORMANCE_EXPERIMENTS.md

@@ -0,0 +1,155 @@
+# EasyMagpie performance experiments
+
+This note records performance experiments that were intentionally not kept in
+the production implementation. Measurements were taken on an RTX A6000 with
+vLLM/vLLM-Omni 0.24 in the `easymagpie-vllm` environment on 2026-07-20 and
+2026-07-21. Treat the absolute numbers as hardware-specific; the controlled
+comparisons are the useful part.
+
+## Native vLLM codec migration
+
+The native codec keeps all 26 causal-convolution histories in vLLM state
+pages, so each call processes only new acoustic frames. The previous exported
+codec replayed the 11-frame left context on every call. In the controlled FP32
+standalone A6000 benchmark, the native path was 1.78x faster at batch 1 and
+2.67x faster at batch 32 for the deployed steady shape (8 new frames versus a
+19-frame stateless replay). Across batch sizes 1–32, the speedup was 1.78–2.67x
+for 8 new frames and 1.80–3.73x for 4 new frames.
+
+The immediately adjacent 128-request, c=32 service measurements (pre-native
+`efd36a9da`, native `1f2745b16`) recorded the following end-to-end change:
+
+| Input | Exported codec | Native vLLM codec | Change |
+| --- | --- | --- | --- |
+| Whole text | 8.30 req/s, 44.08x RTF, 630.2 ms TTFA | 8.73 req/s, 51.79x RTF, 512.9 ms TTFA | +5.2% req/s, +17.5% RTF, -117.3 ms TTFA |
+| Incremental | 7.03 req/s, 40.47x RTF, 723.8 ms TTFA | 7.67 req/s, 46.98x RTF, 614.5 ms TTFA | +9.1% req/s, +16.1% RTF, -109.3 ms TTFA |
+
+Both before/after runs had zero playback underruns. These service numbers
+measure the complete migration rather than codec kernels alone: the native run
+also used FP16 instead of the exported codec's FP32 and changed startup cadence
+from 6/6/8 to 5/6/8 frames. The isolated FP32 results above are therefore the
+better estimate of codec implementation efficiency; the service table captures
+the observed deployment-level benefit at the time. The current deployment uses
+native FP32 because later listening tests found audible artifacts with FP16.
+
+## Codec CUDA graphs (discarded)
+
+The experiment added a codec-specific vLLM runner, a persistent acoustic-code
+input buffer, and exact-shape CUDA graphs. It captured uniform batches for
+chunk lengths 5, 6, and 8 frames and batch sizes 1, 2, 4, 8, 16, 24, and 32.
+Mixed batches, final tails, and shapes outside that set stayed eager. Requests
+were never padded because dummy rows would update the stateful codec cache.
+
+The implementation required about 185 lines in the custom runner and about 291
+production lines overall. It captured 21 graphs in roughly one second and used
+about 0.56 GiB of extra GPU memory. Temporary dispatch instrumentation found
+that only 318 of 592 Stage-1 executions (53.7%) used a graph; the rest were
+mixed or otherwise unmatched and ran eager.
+
+### Results
+
+| Input / concurrency | Eager baseline | Codec CUDA graphs |
+| --- | --- | --- |
+| Whole text, c=1 | 1.20 req/s, 7.33x RTF, 127.8 ms TTFA, 0/1,350 underruns | 1.14 req/s, 7.34x RTF, 127.7 ms TTFA, 0/1,424 underruns |
+| Incremental, c=1 | 1.14 req/s, 6.85x RTF, 134.0 ms TTFA, 0/1,286 underruns | 1.18 req/s, 6.86x RTF, 132.8 ms TTFA, 0/1,243 underruns |
+| Whole text, c=32 | 8.76 req/s, 52.79x RTF, 514.0 ms TTFA, 0/1,337 underruns | 8.70 req/s, 51.08x RTF, 515.4 ms TTFA, 0/1,307 underruns |
+| Incremental, c=32 | 8.02 req/s, 45.90x RTF, 608.3 ms TTFA, 0/1,225 underruns | 7.87 req/s, 44.56x RTF, 613.0 ms TTFA, 2/1,212 underruns |
+
+A second whole-text c=32 graph run produced 8.48 req/s, 53.02x RTF, and
+516.5 ms TTFA. The experiment passed 68 tests (one skipped), produced 128 valid
+mono 22.05 kHz waveforms, and had 1.49% WER (21 substitutions, 6 deletions, and
+2 insertions over 1,941 reference words).
+
+Conclusion: the native codec is already efficient enough that graph launch
+savings were offset by dispatch/copy overhead and incomplete shape coverage.
+The result was neutral to slightly negative, so the custom runner and graph
+machinery were removed. A future H100 experiment could revisit only the steady
+8-frame shape, but it should first demonstrate a material gain in a prototype.
+
+### Codec dtype side result
+
+A controlled 128-request, c=32 Stage-1 comparison favored FP16 over FP32:
+
+| Dtype | RTF (two runs) | Mean TTFA (two runs) | Request rate (two runs) |
+| --- | --- | --- | --- |
+| FP16 | 52.71x / 51.32x | 505.7 / 509.7 ms | 8.77 / 8.51 req/s |
+| FP32 | 45.50x / 44.72x | 608.3 / 608.8 ms | 7.56 / 7.44 req/s |
+
+FP16 improved average RTF by about 15.3% and reduced TTFA by about 101 ms, but
+listening tests exposed audible codec artifacts. The default deployment therefore
+uses FP32; FP16 remains a throughput reference rather than a recommended mode.
+
+## Component no-op ablations
+
+These experiments estimate where end-to-end service time is spent. They are
+compute upper bounds, not candidate model implementations.
+
+### Controlled setup
+
+- Stage 0 used vLLM's dummy weight loader in every variant, including the
+  baseline. Stage 1 kept real codec weights except in the codec-forward no-op.
+- Every request ran exactly 128 Stage-0 decode steps. Audio EOS and all other
+  special acoustic tokens were blocked, and engine stop IDs were disabled.
+- The local-transformer no-op returned a fixed valid code (`1`) for all 16
+  codebooks. The one-pass variant ran the 3-layer local transformer once and
+  then ran all 16 output heads/top-k samplers; it removed only the other 15
+  transformer forwards. The codec no-op returned correctly sized zero audio.
+- Input selection was deterministic (`random.seed(0)`). Concurrency 1 used 32
+  measured requests; concurrency 32 used 128. The benchmark's normal warm-up
+  ran before each measured level.
+- All variants used the same 5/6/8-frame codec streaming cadence. No run had a
+  playback underrun.
+
+The dummy loader skips the model's custom `load_weights`, so the benchmark had
+to initialize the forbidden acoustic-token mask in the constructor. A first
+attempt also used `min_tokens=128`; that exposed the checkpoint's intentionally
+out-of-vocabulary backbone EOS ID (`2` for a two-token dummy vocabulary) in
+vLLM's min-token mask and caused a CUDA index assertion. The final harness used
+`ignore_eos`, no stop IDs, and `max_tokens=128` instead.
+
+### Results
+
+| Variant | c=1: req/s | c=1: RTF | c=1: TTFA | c=32: req/s | c=32: RTF | c=32: TTFA |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: |
+| Baseline: 16 LT forwards + real codec | 0.77 | 7.96x | 127.8 ms | 7.03 / 7.29 | 70.71x / 73.00x | 421.7 / 414.3 ms |
+| One LT forward + real codec | 1.71 | 17.60x | 71.3 ms | 12.72 / 13.52 | 129.42x / 133.39x | 574.9 / 388.1 ms |
+| LT no-op + real codec | 2.10 | 21.68x | 60.7 ms | 16.51 / 17.00 | 165.27x / 170.95x | 202.9 / 208.3 ms |
+| 16 LT forwards + codec no-op | 0.81 | 8.32x | 124.0 ms | 7.94 / 7.96 | 79.75x / 79.63x | 382.1 / 371.7 ms |
+
+Additional mean values for the two c=32 runs:
+
+| Variant | req/s | RTF | TTFA | ITL | Request latency |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| Baseline | 7.16 | 71.86x | 418.0 ms | 262.0 ms | 4.45 s |
+| One LT forward | 13.12 | 131.41x | 481.5 ms | 126.8 ms | 2.43 s |
+| LT no-op | 16.76 | 168.11x | 205.6 ms | 109.9 ms | 1.90 s |
+| Codec no-op | 7.95 | 79.69x | 376.9 ms | 236.6 ms | 4.02 s |
+
+The one-pass TTFA had one high-variance c=32 run (574.9 ms); the repeat was
+388.1 ms. Throughput and end-to-end latency were much more stable and are the
+better signals for component attribution.
+
+### Interpretation
+
+At c=32, using average request rates and comparing service time (`1 / req/s`):
+
+- Removing 15 of the 16 local-transformer forwards reduced baseline service
+  time by about 45% and increased throughput by about 83%.
+- Removing the complete local-transformer/code-prediction path reduced service
+  time by about 57% and increased throughput by about 134%.
+- The difference between the one-pass and full no-op results is about 12% of
+  baseline service time. It includes the remaining transformer forward, all 16
+  output projections, top-k/Gumbel sampling, and staging overhead.
+- Removing codec compute reduced service time by about 10% and increased
+  throughput by about 11%. At c=1, the improvement was only about 5%.
+
+The dominant optimization target is therefore the intra-frame codebook
+predictor, especially repeating the full local transformer 16 times. The codec
+is measurable but secondary after the native vLLM conversion. The ablations do
+not prove that a one-pass predictor preserves quality: it removes autoregressive
+conditioning between codebooks and would require retraining or an architecture
+change before it could become a real optimization.
+
+All hardcoded ablation changes were reverted after measurement. No ablation
+audio is suitable for WER or perceptual evaluation: dummy weights, fixed codes,
+and silent codec output deliberately destroy speech quality.

examples/tts/easymagpie_vllm_omni/README.md

@@ -0,0 +1,114 @@
+## EasyMagpieTTS — vLLM-Omni two-stage inference
+
+Streaming TTS for **NemotronTTS** (Nemotron-H backbone + per-codebook local
+transformer over a 25 fps spectral codec) via [vLLM-Omni](https://github.com/vllm-project/vllm-omni).
+
+The in-engine **two-stage pipeline** chains:
+
+| Stage | Role |
+|-------|------|
+| **0 — talker** | Autoregressive Nemotron-H + local transformer → stacked acoustic codes |
+| **1 — codec** | Stateful native vLLM codec → 22.05 kHz waveform |
+
+Model definition and pipeline registration live in
+[`easymagpie_vllm_omni/`](easymagpie_vllm_omni/) and
+[`vllm_plugin_easymagpie_omni/`](vllm_plugin_easymagpie_omni/).
+Deployment knobs are in [`deploy/easymagpie.yaml`](deploy/easymagpie.yaml).
+
+### Convert a NeMo checkpoint
+
+This step turns the training-time `.nemo` checkpoints into a self-contained
+vLLM-Omni model directory: it converts the talker and causal codec to native
+vLLM models, precomputes the text-embedding lookup, and saves the tokenizer and
+optional speaker embedding. Run it in the **NeMo environment** from the repository root:
+
+```bash
+python examples/tts/easymagpie_vllm_omni/scripts/convert_to_vllm.py \
+  --nemo_file /path/to/emptts.nemo \
+  --codec_model_path /path/to/25fps_spectral_codec.nemo \
+  --phoneme_tokenizer_path /path/to/bpe_ipa_tokenizer.json \
+  --outdir examples/tts/easymagpie_vllm_omni/converted_model \
+  --context_audio /path/to/reference_voice.wav \
+  --speaker_name eng
+```
+
+### Setup the serving environment
+
+Serving needs a GPU, matching **vLLM 0.24 / vLLM-Omni 0.24** versions, and this package.
+It does not need NeMo after conversion:
+
+```bash
+cd examples/tts/easymagpie_vllm_omni
+conda create -n easymagpie-vllm python=3.12 -y
+conda activate easymagpie-vllm
+pip install -r requirements.txt
+pip install -e .
+# optionally for notebook
+pip install ipykernel
+python -m ipykernel install --user \
+  --name easymagpie-vllm \
+  --display-name "Python (easymagpie-vllm)"
+```
+
+Mamba's selective-state-update kernel requires shape- and GPU-specific tuning, so an untuned cache can give
+suboptimal performance. Reuse the same Triton/vLLM cache directories across launches so repeated runs accumulate
+better kernels; for an explicit sweep, run `python scripts/tune_mamba_ssu.py --model converted_model` and restart.
+
+### Quick start — offline synthesis
+
+See [`scripts/offline_demo.ipynb`](scripts/offline_demo.ipynb) to check how `AsyncOmni` is initialized and used.
+
+### Serve over HTTP and WebSocket
+
+```bash
+bash ./scripts/run_server.sh ./converted_model 8091
+```
+
+This starts `vllm serve` with the EasyMagpie plugin on port 8091. Two serving
+APIs are available:
+
+- `POST /v1/audio/speech` with a complete text input.
+- `WS /v1/audio/speech/stream` with incremental text/token updates and
+  asynchronous PCM audio output.
+
+Query the HTTP endpoint from any OpenAI-compatible client:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+  -H 'Content-Type: application/json' \
+  -d '{"input":"This is a TTS service test.","voice":"eng","response_format":"wav","stream":true,"stream_format":"audio"}' \
+  --output out.wav
+```
+
+See [`scripts/server_request.ipynb`](scripts/server_request.ipynb) for examples
+of both serving APIs.
+
+### Benchmarks
+
+```bash
+# Benchmark acoustic token prediction only (no codec).
+python scripts/benchmark_model.py --model ./converted_model -n 128 -c 1 32 \
+    [--incremental --tokens-per-chunk 5]
+
+# Benchmark the service's HTTP API.
+python scripts/benchmark_server.py --text-file vctk_subset.txt -n 128 -c 1 32
+
+# Benchmark the service's incremental synthesis via its WebSocket API.
+python scripts/benchmark_incremental_server.py --model ./converted_model \
+    --text-file vctk_subset.txt --tokens-per-chunk 5 -n 128 -c 1 32
+```
+
+#### RTX A6000 reference (2026-07-21)
+
+Results for 128 requests using VCTK texts. Service rows use the default FP32 codec with 6/6/8-frame cadence:
+
+| Concurrency | Benchmark | Input | Requests/s | RTF | Mean latency | Underruns |
+| ---: | --- | --- | ---: | ---: | ---: | ---: |
+| 1 | Model | Whole text | 1.32 | 7.85x | 28.7 ms TTFT | — |
+| 1 | Model | 5 tokens/chunk | 1.28 | 7.64x | 29.3 ms TTFT | — |
+| 1 | HTTP service | Whole text | 1.27 | 7.73x | 130.8 ms TTFA | 0 / 1,335 |
+| 1 | WebSocket service | 5 tokens/chunk | 1.25 | 7.24x | 137.2 ms TTFA | 0 / 1,226 |
+| 32 | Model | Whole text | 11.20 | 67.75x | 107.3 ms TTFT | — |
+| 32 | Model | 5 tokens/chunk | 10.44 | 63.01x | 105.1 ms TTFT | — |
+| 32 | HTTP service | Whole text | 7.47 | 48.43x | 606.0 ms TTFA | 0 / 1,426 |
+| 32 | WebSocket service | 5 tokens/chunk | 7.45 | 44.12x | 690.9 ms TTFA | 0 / 1,249 |

examples/tts/easymagpie_vllm_omni/deploy/easymagpie.yaml

@@ -0,0 +1,64 @@
+# Stage 0 produces stacked acoustic codes; stage 1 decodes them to audio.
+# Explicit routing is required because the talker reports model_type "nemotron_h".
+pipeline: easymagpie
+
+# Required for streaming code transfer between the two stages.
+async_chunk: true
+
+trust_remote_code: true
+dtype: float16
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      # Native vLLM state pages retain the decoder's exact causal history.
+      # Two 6-frame startup chunks prevent FP32 playback underruns on the RTX A6000.
+      codec_chunk_frames: 8
+      codec_startup_chunk_frames: [6, 6]
+
+stages:
+  # Keep engine settings flat at the stage level; nested settings are not parsed.
+  - stage_id: 0
+    max_num_seqs: 32
+    gpu_memory_utilization: 0.5
+    enforce_eager: false
+    async_scheduling: true
+    enable_prefix_caching: false
+    enable_chunked_prefill: true
+    max_num_batched_tokens: 4096
+    max_model_len: 4096
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    mamba_ssm_cache_dtype: float32
+    skip_tokenizer_init: true
+    engine_extras:
+      worker_cls: easymagpie_vllm_omni.runner.EasyMagpieGPUARWorker
+    default_sampling_params:
+      temperature: 0.0
+      detokenize: false
+      max_tokens: 2048
+
+  - stage_id: 1
+    max_num_seqs: 32
+    gpu_memory_utilization: 0.35
+    # FP16 is faster, but introduces audible codec artifacts.
+    dtype: float32
+    enforce_eager: true
+    async_scheduling: false
+    enable_prefix_caching: false
+    # Codec chunks are atomic: never split one request's state update across steps.
+    enable_chunked_prefill: false
+    max_num_batched_tokens: 512
+    max_model_len: 512
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 65536

examples/tts/easymagpie_vllm_omni/deploy/easymagpie_talker.yaml

@@ -0,0 +1,34 @@
+# EasyMagpieTTS talker-only deploy: acoustic-token prediction without Code2Wav.
+#
+# Use this topology to benchmark and develop the AR talker in isolation from the
+# two-stage codec path. Select it via:
+#   pipeline: easymagpie_talker
+#
+# scripts/benchmark_model.py uses this deploy config by default.
+pipeline: easymagpie_talker
+
+# Required for streaming-text input (whole-text mode ignores this).
+async_chunk: true
+
+trust_remote_code: true
+dtype: float16
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 32
+    gpu_memory_utilization: 0.5
+    enforce_eager: false
+    async_scheduling: true
+    enable_prefix_caching: false
+    enable_chunked_prefill: true
+    max_num_batched_tokens: 4096
+    max_model_len: 4096
+    devices: "0"
+    mamba_ssm_cache_dtype: float32
+    skip_tokenizer_init: true
+    engine_extras:
+      worker_cls: easymagpie_vllm_omni.runner.EasyMagpieGPUARWorker
+    default_sampling_params:
+      temperature: 0.0
+      detokenize: false
+      max_tokens: 2048