Optimize CUDA inference path

README.md

@@ -94,6 +94,9 @@ Try OmniVoice without coding:
 
 - Launch the local web UI: `omnivoice-demo --ip 0.0.0.0 --port 8001`
 
+- Launch the OpenAI-compatible TTS server: `omnivoice-openai-tts-server --host 0.0.0.0 --port 6655`
+
+- Open the browser UI: `http://127.0.0.1:6655/ui`
 
 - Or try it directly on [HuggingFace Space](https://huggingface.co/spaces/k2-fsa/OmniVoice)
 
@@ -113,13 +116,11 @@ Clone a voice from a short reference audio. Provide `ref_audio` and `ref_text`:
 
 ```python
 from omnivoice import OmniVoice
-import torch
 import torchaudio
 
 model = OmniVoice.from_pretrained(
     "k2-fsa/OmniVoice",
     device_map="cuda:0",
-    dtype=torch.float16
 )
 # Apple Silicon users: use device_map="mps" instead
 
@@ -138,6 +139,7 @@ torchaudio.save("out.wav", audio[0], 24000)
 > **Tips**
 >
 > - Use a 3–10 seconds reference audio clip. Longer audio slows down inference and may degrade cloning quality.
+> - On CUDA, OmniVoice now auto-selects an inference-friendly dtype (bf16 on Ampere-or-newer GPUs when available, otherwise fp16) and enables TF32 matmul fast paths.
 > - For better results with Arabic numerals, normalize them to words first (e.g., "123" → "one hundred twenty-three") with text normalization tools (e.g., [WeTextProcessing](https://github.com/wenet-e2e/WeTextProcessing)).
 
 ### Voice Design
@@ -216,6 +218,7 @@ Three CLI entry points are provided. The CLI tools support all features availabl
 | `omnivoice-demo` | Interactive Gradio web demo | [omnivoice/cli/demo.py](omnivoice/cli/demo.py) |
 | `omnivoice-infer` | Single-item inference | [omnivoice/cli/infer.py](omnivoice/cli/infer.py) |
 | `omnivoice-infer-batch` | Batch inference across multiple GPUs | [omnivoice/cli/infer_batch.py](omnivoice/cli/infer_batch.py) |
+| `omnivoice-openai-tts-server` | OpenAI-compatible FastAPI TTS server with request sanitization and sentence-aware chunking | [omnivoice/openai_tts_server.py](omnivoice/openai_tts_server.py) |
 
 ### Demo
 
@@ -269,6 +272,28 @@ Only `id` and `text` are mandatory fields. `ref_audio` and `ref_text` are used i
 
 `language_id`, `language_name`, `duration`, and `speed` are optional. `duration` (in seconds) fixes the output length; `speed` controls the speaking rate. If `duration` and `speed` are both provided, `speed` will be ignored.
 
+### OpenAI-Compatible TTS Server
+
+```bash
+omnivoice-openai-tts-server --host 0.0.0.0 --port 6655
+```
+
+The server exposes `/v1/audio/speech`, `/v1/audio/voices`, `/v1/audio/models`, and `/health` endpoints. Incoming text is sanitized before synthesis (control-token stripping, whitespace cleanup, optional URL/email/phone normalization), and longer inputs are chunked on sentence punctuation before the final audio is returned as one response.
+
+The `/ui` page provides a small browser-friendly credits and voice summary view.
+
+```bash
+curl -X POST http://127.0.0.1:6655/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "omnivoice",
+    "voice": "nova",
+    "input": "Hello world. This sentence-aware request will be sanitized and synthesized.",
+    "response_format": "mp3"
+  }' \
+  --output speech.mp3
+```
+
 ---
 
 ## Training & Evaluation

docs/inference-performance-audit.md

@@ -0,0 +1,34 @@
+# OmniVoice inference performance audit
+
+This note records the current inference-path bottlenecks that were reviewed with an emphasis on **CUDA Ampere GPUs**, plus the fixes landed in this repository.
+
+## Fixed bottlenecks
+
+| Area | Bottleneck | Change |
+| --- | --- | --- |
+| Iterative decoding attention | `omnivoice/models/omnivoice.py` built a dense `(2B, 1, L, L)` boolean attention mask for every inference batch. That is quadratic in sequence length, wastes VRAM, and prevents the inference path from taking advantage of the flex-attention fast path. | Replaced the dense mask with a **BlockMask** built from per-sequence document ids. Padding tokens now keep the old diagonal-only behavior without materializing a dense mask. CUDA inference now defaults to `attn_implementation="flex_attention"` unless the caller overrides it. |
+| Logit scoring | The decoder converted the **entire** batched logits tensor to `float32` before slicing out the per-sample target spans. This added avoidable memory traffic and temporary allocations on every decoding step. | Kept batched logits in model dtype and only promote the smaller per-sample slices inside `_predict_tokens_with_scoring()` where the higher-precision log-softmax is actually needed. |
+| CUDA runtime selection | The single-item CLI, batch CLI, demo, and server each hard-coded their own device logic, and the GPU-facing entry points always loaded the model in `float16`. That missed a safe Ampere optimization path and made runtime behavior inconsistent across tools. | Centralized runtime helpers in `omnivoice.utils.common`: device detection, auto dtype resolution, and TF32 enablement. Inference now auto-selects **bf16 on Ampere-or-newer CUDA GPUs** when available, otherwise fp16 on CUDA and fp32 elsewhere. |
+
+## Remaining bottlenecks worth tracking
+
+These are still present after the changes above and should be treated as follow-up work rather than regressions:
+
+1. **Reference prompt preprocessing is still repeated outside the server cache path.** The OpenAI-compatible server caches `voice_clone_prompt` objects, but the batch CLI still rebuilds prompts if the same reference audio is used across multiple jobs.
+2. **Batch-duration clustering reloads reference audio on CPU.** `omnivoice/cli/infer_batch.py` loads each reference waveform once for duration estimation and again later during generation. A metadata-only duration path would reduce startup cost for large JSONL jobs.
+3. **Chunked long-form generation remains sequential within an utterance.** `omnivoice/models/omnivoice.py::_generate_chunked()` batches across items, but chunk dependencies still serialize work within each long sample.
+4. **Audio decode/post-process work is still per-item.** The final tokenizer decode, silence trimming, cross-fade, and fade/pad steps remain outside the main generation loop and are not batched.
+
+## Ampere-specific guidance
+
+For RTX 30-series and other Ampere GPUs, the best current path in this repo is:
+
+- use the default CUDA device selection
+- let `OmniVoice.from_pretrained(..., device_map="cuda")` choose the dtype automatically
+- keep `nj_per_gpu=1` unless you have enough VRAM to hold multiple full model copies per card
+
+That combination now gives the inference path:
+
+- BF16 weights/activations when the GPU supports them
+- TF32 matmul fast paths for remaining FP32 kernels
+- flex-attention block masks instead of dense quadratic masks

omnivoice/cli/demo.py

@@ -29,21 +29,12 @@
 
 import gradio as gr
 import numpy as np
-import torch
 
 from omnivoice import OmniVoice, OmniVoiceGenerationConfig
+from omnivoice.utils.common import get_best_device
 from omnivoice.utils.lang_map import LANG_NAMES, lang_display_name
 
 
-def get_best_device():
-    """Auto-detect the best available device: CUDA > MPS > CPU."""
-    if torch.cuda.is_available():
-        return "cuda"
-    if torch.backends.mps.is_available():
-        return "mps"
-    return "cpu"
-
-
 # ---------------------------------------------------------------------------
 # Language list — all 600+ supported languages
 # ---------------------------------------------------------------------------
@@ -522,7 +513,6 @@ def main(argv=None) -> int:
     model = OmniVoice.from_pretrained(
         checkpoint,
         device_map=device,
-        dtype=torch.float16,
         load_asr=not args.no_asr,
     )
     print("Model loaded.")

omnivoice/cli/infer.py

@@ -22,20 +22,10 @@
 import argparse
 import logging
 
-import torch
 import torchaudio
 
 from omnivoice.models.omnivoice import OmniVoice
-from omnivoice.utils.common import str2bool
-
-
-def get_best_device():
-    """Auto-detect the best available device: CUDA > MPS > CPU."""
-    if torch.cuda.is_available():
-        return "cuda"
-    if torch.backends.mps.is_available():
-        return "mps"
-    return "cpu"
+from omnivoice.utils.common import get_best_device, str2bool
 
 
 def get_parser() -> argparse.ArgumentParser:
@@ -126,9 +116,7 @@ def main():
 
     device = args.device or get_best_device()
     logging.info(f"Loading model from {args.model} on {device} ...")
-    model = OmniVoice.from_pretrained(
-        args.model, device_map=device, dtype=torch.float16
-    )
+    model = OmniVoice.from_pretrained(args.model, device_map=device)
 
     logging.info(f"Generating audio for: {args.text[:80]}...")
     audios = model.generate(

omnivoice/cli/infer_batch.py

@@ -47,20 +47,11 @@
 
 from omnivoice.models.omnivoice import OmniVoice
 from omnivoice.utils.audio import load_audio
-from omnivoice.utils.common import str2bool
+from omnivoice.utils.common import get_best_device_and_count, str2bool
 from omnivoice.utils.data_utils import read_test_list
 from omnivoice.utils.duration import RuleDurationEstimator
 
 
-def get_best_device():
-    """Auto-detect the best available device: CUDA > MPS > CPU."""
-    if torch.cuda.is_available():
-        return "cuda", torch.cuda.device_count()
-    if torch.backends.mps.is_available():
-        return "mps", 1
-    return "cpu", 1
-
-
 worker_model = None
 SAMPLING_RATE = 24000
 
@@ -231,7 +222,6 @@ def process_init(rank_queue, model_checkpoint, warmup=0):
     worker_model = OmniVoice.from_pretrained(
         model_checkpoint,
         device_map=worker_device,
-        dtype=torch.float16,
     )
 
     if warmup > 0:
@@ -407,11 +397,18 @@ def main():
     args = get_parser().parse_args()
     os.makedirs(args.res_dir, exist_ok=True)
 
-    device_type, num_devices = get_best_device()
+    device_type, num_devices = get_best_device_and_count()
     if device_type == "cpu":
         logging.warning(
             "No GPU found. Falling back to CPU inference. This might be slow."
         )
+    elif device_type.startswith("cuda") and args.nj_per_gpu > 1:
+        logging.warning(
+            "nj_per_gpu=%d will load the full model %d time(s) per GPU. "
+            "For large checkpoints this usually hurts VRAM efficiency and latency.",
+            args.nj_per_gpu,
+            args.nj_per_gpu,
+        )
 
     num_processes = num_devices * args.nj_per_gpu
     logging.info(