[RFC]: General TTS Model Implementation

Motivation.

vLLM-Omni aims to provide fast and easy-to-use serving for TTS models. With growing demand for TTS support, we have multiple models in progress or planned:

TTS Models in vLLM-Omni:

Model	Vendor	Issue	PR	Status
Qwen3-TTS	Alibaba	#938	#1161	✅ Merged
CosyVoice3	Alibaba	#372	#498	🔄 WIP
Step-Audio2	StepFun	#271	#464	🔄 WIP
VoxCPM	-	#1154	-	📋 Open
MiMo-Audio	Xiaomi	#151	#750	✅ Merged
MOSS-TTS	OpenMOSS	#1319	-	📋 Open
Qwen3-Omni (audio)	Alibaba	-	-	✅ Merged
Ming-omni-tts	Ant Group	#1461	-	📋 Open

To ensure consistency and code reuse across these models, new TTS implementations should align with Qwen3-Omni's existing infrastructure, particularly the async_chunk framework for inter-stage streaming.

Key Insight: Qwen3-Omni already has a mature async_chunk framework (PR #1151, #727) that provides inter-stage streaming with 66% latency reduction. TTS models should reuse this framework rather than creating TTS-specific abstractions.

Architecture Overview:

Model	Architecture	Stages
Qwen3-TTS	AR → SpeechTokenizer	2
CosyVoice3	AR → Flow → HiFiGAN	2-3
VoxCPM	(AR↔DIT) → Vocoder	2
Step-Audio2	TBD	TBD
MiMo-Audio	TBD	TBD
MOSS-TTS	AR → AudioDecoder	2
Qwen2.5-Omni	Thinker → Talker → Code2Wav	3
Qwen3-Omni	Thinker → Talker → Code2Wav	3
Ming-omni-tts	(AR + FlowMatching) → AudioVAE	2

Proposed Change.

1 Two-Level Streaming Architecture

TTS models in vLLM-Omni should support two levels of streaming:

┌─────────────────────────────────────────────────────────────────────────┐
│                       Two-Level Streaming Architecture                   │
├─────────────────────────────────────────────────────────────────────────┤
│                                                                          │
│   Level 1: Inter-Stage Streaming (async_chunk)                          │
│   ├─ Enabled via YAML: async_chunk: true                                │
│   ├─ Uses connector for inter-stage data transfer                       │
│   └─ Reuses Qwen3-Omni's stage_input_processor pattern                  │
│                                                                          │
│   Level 2: Output Streaming (AsyncDecodingPipeline)                     │
│   ├─ Background thread for parallel decode                              │
│   ├─ Low first-chunk latency for user experience                        │
│   └─ chunk_size=25, left_context_size=25 (unofficial defaults)           │
│                                                                          │
└─────────────────────────────────────────────────────────────────────────┘

Data Flow Example (Qwen3-TTS):

                    Level 1: async_chunk                Level 2: Output Streaming
                    (inter-stage)                       (to user)
                          │                                   │
                          ▼                                   ▼
┌─────────────────┐     chunk      ┌─────────────────┐     audio     ┌──────────┐
│  Stage 0        │ ─────────────► │  Stage 1        │ ─────────────►│  User    │
│  Talker (AR)    │   connector    │  SpeechTokenizer│   SSE/WS      │  Client  │
│                 │                │  (Code2Wav)     │               │          │
└─────────────────┘                └─────────────────┘               └──────────┘
     │                                   │
     │ generate codec tokens             │ decode to audio
     │ every 25 tokens                   │ in background thread
     └─► put via connector               └─► yield audio chunks

2 Stage Configuration (YAML)

All TTS models should follow the same YAML pattern as qwen3_omni_moe_async_chunk.yaml.

2.1 Required Fields

# stage_configs/<model>_async_chunk.yaml

async_chunk: true  # Enable inter-stage streaming

stage_args:
  - stage_id: 0
    stage_type: llm
    runtime:
      devices: "0"
      max_batch_size: 64
    engine_args:
      model_stage: <stage_name>           # e.g., "talker", "ar"
      model_arch: <ModelClass>            # e.g., Qwen3TTSForConditionalGeneration
      worker_type: ar                     # "ar" for autoregressive, "generation" for decoder
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
      engine_output_type: latent          # "latent" for intermediate, "audio" for final
      custom_process_next_stage_input_func: <module.function>  # Stage input processor
    default_sampling_params:
      temperature: 0.9
      top_k: 50
      max_tokens: 2048

  - stage_id: 1
    stage_type: llm
    runtime:
      devices: "0"  # Can be different GPU
    engine_args:
      model_stage: <decoder_stage>        # e.g., "speech_tokenizer", "code2wav"
      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      engine_output_type: audio
    engine_input_source: [0]              # Receives input from stage 0
    final_output: true
    final_output_type: audio

2.2 Model-Specific Examples

Qwen3-TTS (2-stage):

# stage_configs/qwen3_tts_async_chunk.yaml
async_chunk: true
stage_args:
  - stage_id: 0
    engine_args:
      model_stage: talker
      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2speechtokenizer_async_chunk
  - stage_id: 1
    engine_args:
      model_stage: speech_tokenizer
    engine_input_source: [0]
    final_output: true
    final_output_type: audio

CosyVoice3 (2-stage):

# stage_configs/cosyvoice3_async_chunk.yaml
async_chunk: true
stage_args:
  - stage_id: 0
    engine_args:
      model_stage: ar
      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.cosyvoice3.ar2vocoder_async_chunk
  - stage_id: 1
    engine_args:
      model_stage: flow_hifigan
    engine_input_source: [0]
    final_output: true
    final_output_type: audio

3 Stage Input Processors

Stage input processors handle data transfer between stages. They follow the pattern established in qwen3_omni.py.

3.1 Function Signature

# vllm_omni/model_executor/stage_input_processors/<model>.py

def <stage_a>2<stage_b>_async_chunk(
    connector: Any,
    pooling_output: dict[str, Any],
    request: OmniEngineCoreRequest,
) -> dict[str, Any] | None:
    """
    Process Stage A outputs to create Stage B inputs.

    Args:
        connector: Inter-stage connector with buffering
        pooling_output: Output from current stage (contains codec codes, hidden states, etc.)
        request: Current request with metadata

    Returns:
        dict with stage B inputs, or None if chunk not ready
    """
    pass

3.2 Naming Convention

Pattern	Example
<source>2<target>_async_chunk	talker2code2wav_async_chunk
<source>2<target> (non-streaming)	talker2code2wav

3.3 Reference Implementation (Qwen3-Omni)

# From qwen3_omni.py - talker2code2wav_async_chunk

def talker2code2wav_async_chunk(
    connector: Any,
    pooling_output: dict[str, Any],
    request: OmniEngineCoreRequest,
):
    if "code_predictor_codes" not in pooling_output:
        return None

    code_predictor_codes = pooling_output["code_predictor_codes"]
    if code_predictor_codes is None or code_predictor_codes.numel() == 0:
        return None

    request_id = request.external_req_id
    chunk_size = left_context_size = 25  # Validated defaults

    # Buffer codes in connector
    codec_codes = code_predictor_codes.to(torch.long).transpose(0, 1).cpu().reshape(-1).tolist()
    connector.code_prompt_token_ids[request_id].append(codec_codes)

    length = len(connector.code_prompt_token_ids[request_id])
    chunk_length = length % chunk_size

    # Only send when chunk is full or finished
    if chunk_length != 0 and not request.is_finished():
        return None

    context_length = chunk_length if chunk_length != 0 else chunk_size
    end_index = min(length, left_context_size + context_length)

    return {
        "code_predictor_codes": (
            torch.tensor(connector.code_prompt_token_ids[request_id][-end_index:])
            .transpose(0, 1).reshape(-1).tolist()
        ),
        "finished": torch.tensor(request.is_finished(), dtype=torch.bool),
    }

3.4 Key Components

Component	Purpose	Usage
connector.code_prompt_token_ids[request_id]	Buffer for codec codes	Accumulate tokens until chunk_size
connector.put_requests[request_id]	Chunk ID counter	Track which chunk we're on
connector.request_payload[request_id]	Temporary storage	Store partial data across calls
chunk_size = 25	Tokens per chunk	Validated in Qwen3-Omni
left_context_size = 25	Context for boundary smoothing	Prevents audio artifacts

4 Chunked Decode Streaming

The final stage (decoder) should implement chunked_decode_streaming() for memory-efficient decoding.

4.1 Method Signature

def chunked_decode_streaming(
    self,
    codes: torch.Tensor,
    chunk_size: int = 25,
    left_context_size: int = 25,
) -> torch.Tensor:
    """
    Decode codec codes to audio with streaming support.

    Args:
        codes: Codec codes from AR stage
        chunk_size: Number of codec frames per chunk
        left_context_size: Overlapping frames for context

    Returns:
        Audio waveform tensor
    """

4.2 Reference Implementation (Qwen3-Omni Code2Wav)

# From qwen3_omni_code2wav.py

def chunked_decode_streaming(
    self,
    codes: torch.Tensor,
    chunk_size: int = 25,
    left_context_size: int = 25,
) -> torch.Tensor:
    wavs = []
    end_index = codes.shape[-1]

    # TODO: need to optimize algorithms, current only support
    # chunk_size = left_context_size = 25
    if end_index <= chunk_size:
        context_size = 0
    else:
        context_size = left_context_size

    wav_chunk = self(codes)
    # Remove context from output
    wavs.append(wav_chunk[..., context_size * self.total_upsample :])
    return torch.cat(wavs, dim=-1)

5 Output Streaming (Level 2)

For low first-chunk latency, implement background thread decoding. The following are proposed abstractions (actual naming TBD with PR #1189).

5.1 StreamingChunkOutput (Proposed)

@dataclass
class StreamingChunkOutput:
    """Output chunk from streaming TTS generation."""
    chunk_idx: int                    # Sequential chunk index (0-based)
    is_finished: bool                 # Whether generation is complete
    total_generated: int              # Total tokens generated so far
    codec_codes: torch.Tensor | None  # Intermediate codec codes (optional)
    audio: np.ndarray | None          # Decoded audio samples
    sample_rate: int = 24000

5.2 Async Decoder (Proposed)

class AsyncDecodingPipeline:
    """Background thread for parallel audio decoding."""

    def __init__(self, decode_fn: Callable, max_queue_size: int = 10):
        self.decode_fn = decode_fn
        self.input_queue = queue.Queue(maxsize=max_queue_size)
        self.output_queue = queue.Queue()
        self.thread = None

    def start(self):
        self.thread = threading.Thread(target=self._decode_loop, daemon=True)
        self.thread.start()

    def submit(self, codec_codes: torch.Tensor, is_last: bool = False):
        self.input_queue.put((codec_codes, is_last))

    def get_result(self, timeout: float | None = None) -> tuple:
        return self.output_queue.get(timeout=timeout)

    def stop(self):
        self.input_queue.put((None, True))  # Sentinel
        if self.thread:
            self.thread.join()

6 Consistency Table

Aspect	Qwen3-Omni	Qwen2.5-Omni	Qwen3-TTS	CosyVoice3	Notes
Level 1: async_chunk	✅ Merged	❌	✅ Merged	❌	Inter-stage streaming
Audio chunk yield	✅ via async_chunk	❌	✅ Merged	❌	Yield audio as generated
Background thread decode	❌	❌	🔄 #1189	❌	Further latency reduction
chunk_size default	25	TBD	25	TBD	Validated value
left_context_size default	25	TBD	25	TBD	Prevents artifacts
Stage processor naming	*_async_chunk	*_async_chunk	*_async_chunk	*_async_chunk	Consistent pattern
Connector usage	code_prompt_token_ids	Same	Same	Same	Reuse infrastructure
chunked_decode_streaming()	✅	❌	✅	❌	Memory efficient

7 New TTS Model Checklist

When adding a new TTS model, follow these steps:

Step 1: Define Stage Architecture

[ ] Identify stages (AR, Flow, Vocoder, etc.)
[ ] Determine stage boundaries (where to split)
[ ] Identify intermediate representations (codec codes, mel, etc.)

Step 2: Create Stage Config YAML

[ ] Create stage_configs/<model>_async_chunk.yaml
[ ] Set async_chunk: true
[ ] Define stage_args with correct model_stage, model_arch
[ ] Set custom_process_next_stage_input_func
[ ] Configure engine_input_source for downstream stages

Step 3: Implement Stage Input Processor

[ ] Create stage_input_processors/<model>.py
[ ] Implement <stage_a>2<stage_b>_async_chunk function
[ ] Use connector.code_prompt_token_ids for buffering
[ ] Use chunk_size=25, left_context_size=25
[ ] Return None if chunk not ready

Step 4: Add chunked_decode_streaming()

[ ] Add method to final stage model
[ ] Use same chunk_size/left_context_size defaults
[ ] Handle context removal for smooth boundaries

Step 5: (Optional) Add AsyncDecodingPipeline

[ ] For lower latency, add background decoding
[ ] Use StreamingChunkOutput for chunk output
[ ] Integrate with API layer (/v1/audio/speech)

Example: Stage Architecture Decisions

Model	Stage 0	Stage 1	Notes
Qwen3-TTS	Talker (AR)	SpeechTokenizer	Simple 2-stage
CosyVoice3	AR	Flow + HiFiGAN	Flow can be separate stage
VoxCPM	AR↔DIT (internal loop)	Vocoder	Internal loops stay in one stage
MOSS-TTS	AR (Delay/Local)	AudioDecoder	2 variants: MossTTSDelay (8B), MossTTSLocal (1.7B)
Qwen2.5-Omni	Thinker → Talker	Code2Wav	3-stage, similar to Qwen3-Omni
Qwen3-Omni	Thinker → Talker	Code2Wav	3-stage for multimodal
Ming-omni-tts	Latent Predictor (AR + FlowLoss inline)	AudioVAE decoder	stop_head replaces EOS; FlowMatching runs within Stage 0 forward()

Feedback Period.

No response

CC List.

@hsliuustc0106 @Gaohan123 @gerayking @gcanlin @linyueqian @Sy0307 @tsdocode @xulusjb @zhaotyer @divyanshsinghvi @amy-why-3459 @R2-Y @ywang96

Any Other Things.

1. Where should streaming output classes (e.g., StreamingChunkOutput, async decoder) live?

   • Option A: vllm_omni/outputs.py (alongside OmniRequestOutput)
   • Option B: vllm_omni/streaming/ (new top-level module)
   • Option C: vllm_omni/model_executor/models/streaming/ (near model code)

   Note: Qwen3-Omni already streams audio chunks via async_chunk (yields audio as Code2Wav generates). PR #1189 adds background thread decoding for further latency reduction. A shared abstraction could benefit both models.

2. Should chunk_size and left_context_size be configurable via /v1/audio/speech API?

3. How does CUDA Graph (#1205) interact with async_chunk? (Fixed input sizes vs variable chunks)

4. Should we formalize the connector interface as a Protocol class?

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.