From fd51841230afb1e21bf6555ca6b439313ebfb4c7 Mon Sep 17 00:00:00 2001 From: Du Bin Date: Fri, 6 Mar 2026 12:22:01 +0800 Subject: [PATCH 1/9] perf: replace per-element .item() GPU syncs with batch .tolist() in TTS code predictor (#1614) Signed-off-by: dubin555 Signed-off-by: linyueqian Co-authored-by: linyueqian Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .../qwen3_tts/qwen3_tts_code_predictor_vllm.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index 4820aa2a7..d78e16b9c 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -120,11 +120,17 @@ class _LocalPredictorKVCache: if num_tokens <= 0: return {}, torch.empty((0,), dtype=torch.int64, device=self.device), {} + # Convert tensors to Python lists once to avoid per-element .item() + # calls (each .item() on a GPU tensor forces a GPU→CPU sync). + query_lens_list = query_lens_i32[:num_reqs].tolist() + seq_lens_list = seq_lens_i32[:num_reqs].tolist() + block_table_cpu = self._block_table[:num_reqs].cpu().tolist() + # positions: for each request i, emit positions [seq_len-query_len .. seq_len-1] pos_list: list[torch.Tensor] = [] for i in range(num_reqs): - ql = int(query_lens_i32[i].item()) - sl = int(seq_lens_i32[i].item()) + ql = query_lens_list[i] + sl = seq_lens_list[i] start = sl - ql pos_list.append(torch.arange(start, sl, dtype=torch.int64)) positions_cpu = torch.cat(pos_list, dim=0) @@ -134,13 +140,13 @@ class _LocalPredictorKVCache: slot_mapping = torch.empty((num_tokens,), dtype=torch.int64, device="cpu") cursor = 0 for i in range(num_reqs): - ql = int(query_lens_i32[i].item()) - sl = int(seq_lens_i32[i].item()) + ql = query_lens_list[i] + sl = seq_lens_list[i] start = sl - ql for p in range(start, sl): block_idx = p // self.block_size offset = p % self.block_size - block_id = int(self._block_table[i, block_idx].item()) + block_id = block_table_cpu[i][block_idx] slot_mapping[cursor] = block_id * self.block_size + offset cursor += 1 From 90ec499bd557537088cbdf50259453b3a2f2d13e Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Fri, 6 Mar 2026 13:21:55 +0800 Subject: [PATCH 2/9] [Refactor][Perf] Qwen3-TTS: re-prefill Code Predictor with torch.compile + enable Code2Wav decoder CUDA Graph (#1617) Signed-off-by: Sy03 <1370724210@qq.com> --- .../models/qwen3_tts/qwen3_tts_code2wav.py | 133 ++-- .../qwen3_tts_code_predictor_vllm.py | 613 +++++++++--------- 2 files changed, 368 insertions(+), 378 deletions(-) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index c93d027d7..eee751e16 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -4,7 +4,6 @@ import os from collections.abc import Iterable from typing import Any -import numpy as np import torch import torch.nn as nn from transformers.utils.hub import cached_file @@ -21,7 +20,8 @@ logger = init_logger(__name__) class Qwen3TTSCode2Wav(nn.Module): """Stage-1 code2wav model for Qwen3-TTS (GenerationModelRunner). - Consumes frame-aligned codec tokens from input_ids and decodes waveform via SpeechTokenizer.""" + Consumes frame-aligned codec tokens from input_ids and decodes waveform + via the SpeechTokenizer decoder directly (bypassing HF wrapper overhead).""" input_modalities = "audio" @@ -34,13 +34,13 @@ class Qwen3TTSCode2Wav(nn.Module): self.has_preprocess = False self.has_postprocess = False self.enable_update_additional_information = True - # Generation-only stage (no logits / sampling). self.requires_raw_input_tokens = True self._speech_tokenizer: Qwen3TTSTokenizer | None = None + self._decoder: nn.Module | None = None self._num_quantizers: int | None = None - self._decode_upsample_rate: int | None = None self._output_sample_rate: int | None = None + self._total_upsample: int | None = None self._logged_codec_stats = False @staticmethod @@ -52,18 +52,15 @@ class Qwen3TTSCode2Wav(nn.Module): return buf.device return torch.device("cpu") - def _ensure_speech_tokenizer_loaded(self) -> Qwen3TTSTokenizer: - if self._speech_tokenizer is not None: - return self._speech_tokenizer + def _ensure_speech_tokenizer_loaded(self) -> None: + if self._decoder is not None: + return - # Locate speech_tokenizer dir from HF cache (or local path). cfg_path = cached_file(self.model_path, "speech_tokenizer/config.json") if cfg_path is None: raise ValueError(f"{self.model_path}/speech_tokenizer/config.json not found") speech_tokenizer_dir = os.path.dirname(cfg_path) - # Stage-1 only needs decode; skip HF feature extractor to avoid heavy optional deps. - # Still require preprocessor_config.json (use cached_file so online runs can fetch it). prep_cfg = cached_file(self.model_path, "speech_tokenizer/preprocessor_config.json") if prep_cfg is None: raise ValueError( @@ -77,12 +74,10 @@ class Qwen3TTSCode2Wav(nn.Module): load_feature_extractor=False, ) - # Align device with vLLM worker, then read back from module. if tok.model is not None: tok.model.to(device=self.vllm_config.device_config.device) tok.device = self._module_device(tok.model) - # Derive codec group count and rates from tokenizer config. dec_cfg = getattr(tok.model.config, "decoder_config", None) num_q = getattr(dec_cfg, "num_quantizers", None) if dec_cfg is not None else None if num_q is None: @@ -103,11 +98,39 @@ class Qwen3TTSCode2Wav(nn.Module): except Exception as e: raise ValueError(f"Failed to get output sample rate: {e}") from e + decoder = tok.model.decoder + decoder.eval() + self._speech_tokenizer = tok + self._decoder = decoder self._num_quantizers = num_q - self._decode_upsample_rate = upsample self._output_sample_rate = out_sr - return tok + self._total_upsample = int(decoder.total_upsample) + + if hasattr(decoder, "enable_cudagraph"): + device = self._module_device(decoder) + if device.type == "cuda": + try: + capture_sizes = None + model_cfg = getattr(self.vllm_config, "model_config", None) + connector_cfg = getattr(model_cfg, "stage_connector_config", None) + extra_cfg = ( + connector_cfg.get("extra", connector_cfg) + if isinstance(connector_cfg, dict) + else getattr(connector_cfg, "extra", None) + ) + if isinstance(extra_cfg, dict): + chunk_frames = int(extra_cfg.get("codec_chunk_frames") or 0) + left_frames = int(extra_cfg.get("codec_left_context_frames") or 0) + if chunk_frames > 0 and left_frames >= 0: + from .cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper + + steady_window = left_frames + chunk_frames + capture_sizes = sorted({*CUDAGraphDecoderWrapper.DEFAULT_CAPTURE_SIZES, steady_window}) + decoder.enable_cudagraph(capture_sizes=capture_sizes, device=device) + logger.info("Code2Wav decoder CUDA Graph enabled") + except Exception: + logger.warning("Failed to enable CUDA Graph for Code2Wav decoder", exc_info=True) def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: # This stage ignores token embeddings. Keep a stable dummy embedding for vLLM runner. @@ -155,17 +178,19 @@ class Qwen3TTSCode2Wav(nn.Module): input_ids layout per request: [codec_context_frames, *flat_codes] where flat_codes is codebook-major [q*F]. - When batched, uses forward context ubatch_slices to split the - concatenated input_ids and decode via a single batched forward pass. + Bypasses the HF Qwen3TTSTokenizer.decode() wrapper and calls the + decoder.chunked_decode() directly to avoid GPU->CPU->GPU round-trips. + Length management is done here instead of relying on HF's padding=-1 + sentinel logic. """ self._ensure_speech_tokenizer_loaded() + assert self._decoder is not None assert self._num_quantizers is not None - assert self._decode_upsample_rate is not None - assert self._output_sample_rate is not None + assert self._total_upsample is not None - tok = self._speech_tokenizer + decoder = self._decoder q = int(self._num_quantizers) - upsample = int(self._decode_upsample_rate) + upsample = int(self._total_upsample) sr_val = int(self._output_sample_rate) sr_tensor = torch.tensor(sr_val, dtype=torch.int32) empty = torch.zeros((0,), dtype=torch.float32) @@ -179,13 +204,9 @@ class Qwen3TTSCode2Wav(nn.Module): ids = input_ids.reshape(-1).to(dtype=torch.long) request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts")) - # Parse each request: extract left_context_size, validate, reshape codes. - # input_ids layout per request: [*flat_codes] - # where flat_codes is codebook-major [q*F]. - parsed = [] # (left_context_size, actual_frames) - valid_codes = [] - valid_indices = [] - # Get left_context_size from runtime_additional_information (passed via kwargs) + parsed: list[tuple[int, int]] = [] + valid_codes_qf: list[torch.Tensor] = [] + valid_indices: list[int] = [] left_context_size = [0] * len(request_ids_list) if runtime_additional_information is not None: for i, info in enumerate(runtime_additional_information): @@ -193,8 +214,6 @@ class Qwen3TTSCode2Wav(nn.Module): break if "left_context_size" in info: left_context_size[i] = info["left_context_size"] - else: - logger.debug("No additional_information provided to code2wav stage.") for i, req_ids in enumerate(request_ids_list): if req_ids.numel() < 1: parsed.append((0, 0)) @@ -202,7 +221,6 @@ class Qwen3TTSCode2Wav(nn.Module): ctx_frames = left_context_size[i] flat = req_ids n = flat.numel() - # Warmup / dummy_run: not divisible by num_quantizers. if n == 0 or n % q != 0: if n > 0: logger.warning( @@ -214,14 +232,14 @@ class Qwen3TTSCode2Wav(nn.Module): parsed.append((0, 0)) continue frames = n // q - # Reshape codebook-major flat [q*F] -> [q, F] -> [F, q] for SpeechTokenizer. - codes_fq = flat.reshape(q, frames).transpose(0, 1).contiguous() + # [q*F] -> [Q, F] for direct decoder call (decoder expects [B, Q, F]) + codes_qf = flat.reshape(q, frames) parsed.append((ctx_frames, frames)) - valid_codes.append({"audio_codes": codes_fq}) + valid_codes_qf.append(codes_qf) valid_indices.append(i) num_req = len(request_ids_list) - if not valid_codes: + if not valid_codes_qf: return OmniOutput( text_hidden_states=None, multimodal_outputs={ @@ -233,50 +251,55 @@ class Qwen3TTSCode2Wav(nn.Module): if not self._logged_codec_stats: self._logged_codec_stats = True try: - c = valid_codes[0]["audio_codes"] + c = valid_codes_qf[0] logger.info( - "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] head=%s batch=%d", - c.shape[0], + "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] batch=%d", + c.shape[1], q, int(torch.unique(c).numel()), int(c.min().item()), int(c.max().item()), - c[: min(2, c.shape[0]), : min(8, q)].cpu().tolist(), - len(valid_codes), + len(valid_codes_qf), ) except Exception: pass - # Batched decode: single forward pass through SpeechTokenizer. - wavs, _ = tok.decode(valid_codes) - if len(wavs) != len(valid_codes): - raise RuntimeError(f"Code2Wav returned {len(wavs)} waveforms for {len(valid_codes)} requests") + # Decode directly via decoder.chunked_decode(), staying entirely on GPU. + # For single request: no padding needed, fast path. + # For multiple requests: decode each individually to avoid padding overhead. + wav_tensors: list[torch.Tensor] = [] + if len(valid_codes_qf) == 1: + codes_bqf = valid_codes_qf[0].unsqueeze(0) # [1, Q, F] + wav = decoder.chunked_decode(codes_bqf) # [1, 1, wav_len] + wav_tensors.append(wav.squeeze(0).squeeze(0)) # [wav_len] + else: + for codes_qf in valid_codes_qf: + codes_bqf = codes_qf.unsqueeze(0) # [1, Q, F] + wav = decoder.chunked_decode(codes_bqf) + wav_tensors.append(wav.squeeze(0).squeeze(0)) - # Build per-request outputs, trimming padding and left-context. - audios = [empty] * num_req + audios: list[torch.Tensor] = [empty] * num_req srs = [sr_tensor] * num_req for j, idx in enumerate(valid_indices): ctx_frames, actual_frames = parsed[idx] - audio_np = wavs[j].astype(np.float32, copy=False) - # Trim decoder padding (output may be longer due to batch padding). + wav = wav_tensors[j] expected_len = actual_frames * upsample - if audio_np.shape[0] > expected_len: - audio_np = audio_np[:expected_len] - # Trim left-context waveform samples (streaming sliding window). + if wav.shape[0] > expected_len: + wav = wav[:expected_len] if ctx_frames > 0: cut = ctx_frames * upsample - if cut < audio_np.shape[0]: - audio_np = audio_np[cut:] + if cut < wav.shape[0]: + wav = wav[cut:] else: logger.warning( "Context trim %d >= decoded length %d; returning empty audio.", cut, - audio_np.shape[0], + wav.shape[0], ) continue - if audio_np.shape[0] > 0: - audios[idx] = torch.from_numpy(audio_np).to(dtype=torch.float32).reshape(-1) + if wav.shape[0] > 0: + audios[idx] = wav.to(dtype=torch.float32).reshape(-1) return OmniOutput( text_hidden_states=None, diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index d78e16b9c..7f3fdfd7a 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -1,235 +1,217 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any import torch import torch.nn as nn +import torch.nn.functional as F from vllm.config import VllmConfig from vllm.config.vllm import set_current_vllm_config -from vllm.forward_context import set_forward_context +from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) -from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer from vllm.model_executor.models.utils import is_pp_missing_parameter -from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, KVCacheTensor -from vllm.v1.worker.gpu import attn_utils - -from vllm_omni.platforms import current_omni_platform from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig +logger = init_logger(__name__) -class _LocalPredictorKVCache: - """Minimal local KV cache + attention metadata for running - code_predictor inside one worker (independent of engine KV).""" + +# =================================================================== +# Standalone Code Predictor Layers (no vLLM paged attention) +# =================================================================== +# +# These replace vLLM's Qwen3DecoderLayer for the code predictor. +# Input is batch-major [B, seq_len, H], attention uses F.scaled_dot_product_attention. +# Weight names match the checkpoint (self_attn.qkv_proj, mlp.gate_up_proj, etc.) +# so load_weights works unchanged. + + +class _CodePredictorAttention(nn.Module): + """Standalone multi-head attention for code predictor. + + Uses F.scaled_dot_product_attention (SDPA) instead of vLLM's paged Attention. + Supports fused QKV, RoPE, q/k normalization, and native GQA via enable_gqa. + Input: [B, seq_len, hidden_size], output: [B, seq_len, hidden_size]. + """ def __init__( self, + config: Qwen3TTSTalkerCodePredictorConfig, *, - vllm_config: VllmConfig, - max_seq_len: int, - max_batch_size: int, - device: torch.device, + prefix: str = "", ) -> None: - self.vllm_config = vllm_config - self.device = device - - # Collect attention layers registered in this vllm_config. - kv_cache_spec_by_layer = attn_utils.get_kv_cache_spec(vllm_config) - if not kv_cache_spec_by_layer: - raise RuntimeError("Local predictor KVCache requires vLLM Attention layers to be registered.") - - # We only need enough blocks for a tiny per-frame sequence (<= max_seq_len). - any_spec = next(iter(kv_cache_spec_by_layer.values())) - block_size = int(any_spec.block_size) - blocks_per_seq = (int(max_seq_len) + block_size - 1) // block_size - num_blocks = max(1, int(max_batch_size) * int(blocks_per_seq)) - - # Allocate per-layer KV caches (small, independent). - kv_cache_tensors: list[KVCacheTensor] = [] - for layer_name, spec in kv_cache_spec_by_layer.items(): - kv_cache_tensors.append(KVCacheTensor(size=int(spec.page_size_bytes) * num_blocks, shared_by=[layer_name])) - - merged_spec: KVCacheSpec = KVCacheSpec.merge(list(kv_cache_spec_by_layer.values())) - self.kv_cache_config = KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=kv_cache_tensors, - kv_cache_groups=[ - KVCacheGroupSpec(layer_names=list(kv_cache_spec_by_layer.keys()), kv_cache_spec=merged_spec) - ], + super().__init__() + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = getattr( + config, + "head_dim", + config.hidden_size // config.num_attention_heads, ) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self._use_gqa = self.num_kv_heads != self.num_heads - # Init backend + bind KV cache tensors to attention modules. - self.attn_backends, self.attn_metadata_builders = attn_utils.init_attn_backend( - self.kv_cache_config, vllm_config, device + self.qkv_proj = QKVParallelLinear( + hidden_size=self.hidden_size, + head_size=self.head_dim, + total_num_heads=self.num_heads, + total_num_kv_heads=self.num_kv_heads, + bias=getattr(config, "attention_bias", False), + prefix=f"{prefix}.qkv_proj", + disable_tp=True, ) - self.runner_kv_caches: list[torch.Tensor] = [] - attn_utils.init_kv_cache( - self.runner_kv_caches, - vllm_config.compilation_config.static_forward_context, - self.kv_cache_config, - self.attn_backends, - device, + self.o_proj = RowParallelLinear( + input_size=self.num_heads * self.head_dim, + output_size=self.hidden_size, + bias=False, + prefix=f"{prefix}.o_proj", + disable_tp=True, ) + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters=getattr(config, "rope_parameters", None), + dual_chunk_attention_config=None, + ) + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) - # Precompute a fixed block table mapping for the maximum batch. - self.block_size = block_size - self.blocks_per_seq = blocks_per_seq - self.max_batch_size = int(max_batch_size) - - bt = torch.full((self.max_batch_size, self.blocks_per_seq), -1, dtype=torch.int32, device=device) - for i in range(self.max_batch_size): - for j in range(self.blocks_per_seq): - bt[i, j] = i * self.blocks_per_seq + j - self._block_table = bt - - def build_attn_metadata( + def forward( self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + bsz, seq_len, _ = hidden_states.shape + + qkv, _ = self.qkv_proj(hidden_states.reshape(bsz * seq_len, -1)) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape) + k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape) + + q, k = self.rotary_emb(position_ids, q, k) + + q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + + attn_out = F.scaled_dot_product_attention( + q, + k, + v, + scale=self.scaling, + is_causal=True, + enable_gqa=self._use_gqa, + ) + + attn_out = attn_out.transpose(1, 2).reshape(bsz * seq_len, -1) + output, _ = self.o_proj(attn_out) + return output.view(bsz, seq_len, -1) + + +class _CodePredictorMLP(nn.Module): + """SiLU-gated MLP for code predictor, matching Qwen3MLP structure.""" + + def __init__( + self, + config: Qwen3TTSTalkerCodePredictorConfig, *, - num_reqs: int, - query_lens: torch.Tensor, # (num_reqs,) int32 on cpu - seq_lens: torch.Tensor, # (num_reqs,) int32 on cpu - ) -> tuple[dict[str, Any], torch.Tensor, dict[str, torch.Tensor]]: - """Build attention metadata, positions, and slot_mapping dict. + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[config.intermediate_size] * 2, + bias=False, + prefix=f"{prefix}.gate_up_proj", + disable_tp=True, + ) + self.down_proj = RowParallelLinear( + input_size=config.intermediate_size, + output_size=config.hidden_size, + bias=False, + prefix=f"{prefix}.down_proj", + disable_tp=True, + ) - Returns: - (attn_metadata, positions, slot_mappings_by_layer) - - attn_metadata: per-layer attention metadata for attn backends. - - positions: (num_tokens,) position IDs on device. - - slot_mappings_by_layer: {layer_name: slot_mapping_tensor} for - set_forward_context so that unified_kv_cache_update can write - the KV cache correctly. - """ - num_reqs = int(num_reqs) - if num_reqs <= 0: - return {}, torch.empty((0,), dtype=torch.int64, device=self.device), {} - if num_reqs > self.max_batch_size: - raise ValueError(f"num_reqs={num_reqs} exceeds local predictor max_batch_size={self.max_batch_size}") + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + gate, up = gate_up.chunk(2, dim=-1) + x = F.silu(gate) * up + x, _ = self.down_proj(x) + return x - query_lens_i32 = query_lens.to(dtype=torch.int32, device="cpu") - seq_lens_i32 = seq_lens.to(dtype=torch.int32, device="cpu") - # query_start_loc: prefix sums of query_lens. - qsl = torch.zeros((num_reqs + 1,), dtype=torch.int32, device="cpu") - qsl[1:] = torch.cumsum(query_lens_i32, dim=0) - num_tokens = int(qsl[-1].item()) - if num_tokens <= 0: - return {}, torch.empty((0,), dtype=torch.int64, device=self.device), {} +class _CodePredictorDecoderLayer(nn.Module): + """Transformer decoder layer for code predictor (SDPA, no KV cache).""" - # Convert tensors to Python lists once to avoid per-element .item() - # calls (each .item() on a GPU tensor forces a GPU→CPU sync). - query_lens_list = query_lens_i32[:num_reqs].tolist() - seq_lens_list = seq_lens_i32[:num_reqs].tolist() - block_table_cpu = self._block_table[:num_reqs].cpu().tolist() + def __init__( + self, + config: Qwen3TTSTalkerCodePredictorConfig, + *, + prefix: str = "", + ) -> None: + super().__init__() + self.self_attn = _CodePredictorAttention(config, prefix=f"{prefix}.self_attn") + self.mlp = _CodePredictorMLP(config, prefix=f"{prefix}.mlp") + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - # positions: for each request i, emit positions [seq_len-query_len .. seq_len-1] - pos_list: list[torch.Tensor] = [] - for i in range(num_reqs): - ql = query_lens_list[i] - sl = seq_lens_list[i] - start = sl - ql - pos_list.append(torch.arange(start, sl, dtype=torch.int64)) - positions_cpu = torch.cat(pos_list, dim=0) + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_ids) + hidden_states = residual + hidden_states - # slot_mapping: map each query token to a physical slot in the paged KV cache. - # We allocate per-request contiguous blocks; slot = base + position. - slot_mapping = torch.empty((num_tokens,), dtype=torch.int64, device="cpu") - cursor = 0 - for i in range(num_reqs): - ql = query_lens_list[i] - sl = seq_lens_list[i] - start = sl - ql - for p in range(start, sl): - block_idx = p // self.block_size - offset = p % self.block_size - block_id = block_table_cpu[i][block_idx] - slot_mapping[cursor] = block_id * self.block_size + offset - cursor += 1 + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states - max_seq_len = int(seq_lens_i32[:num_reqs].max().item()) - query_start_loc_gpu = qsl.to(device=self.device) - seq_lens_gpu = seq_lens_i32.to(device=self.device) - block_table = self._block_table[:num_reqs].contiguous() - slot_mapping_gpu = slot_mapping.to(device=self.device) - # FIXME(gcanlin): Refactor build_attn_metadata to avoid special-casing NPU backends here. - if current_omni_platform.is_npu(): - # NPU requires AscendCommonAttentionMetadata with extra attributes - from vllm_ascend.worker.v2 import attn_utils as attn_utils_npu - - max_query_len = int(query_lens_i32[:num_reqs].max().item()) - # NPU version expects slot_mappings as a stacked tensor, not a list - slot_mappings_tensor = slot_mapping_gpu.unsqueeze(0) - attn_metadata = attn_utils_npu.build_attn_metadata( - attn_metadata_builders=self.attn_metadata_builders, - num_reqs=num_reqs, - num_tokens=num_tokens, - query_start_loc_gpu=query_start_loc_gpu, - query_start_loc_cpu=qsl, - max_query_len=max_query_len, - seq_lens=seq_lens_gpu, - max_seq_len=max_seq_len, - block_tables=[block_table], - slot_mappings=slot_mappings_tensor, - kv_cache_config=self.kv_cache_config, - ) - else: - attn_metadata = attn_utils.build_attn_metadata( - self.attn_metadata_builders, - num_reqs=num_reqs, - num_tokens=num_tokens, - query_start_loc_gpu=query_start_loc_gpu, - query_start_loc_cpu=qsl, - seq_lens=seq_lens_gpu, - max_seq_len=max_seq_len, - block_tables=[block_table], - slot_mappings=[slot_mapping_gpu], - kv_cache_config=self.kv_cache_config, - ) - - # Build slot_mappings_by_layer for set_forward_context. - # Fix for vllm 0.15.0 - slot_mappings_by_layer: dict[str, torch.Tensor] = {} - for kv_cache_group in self.kv_cache_config.kv_cache_groups: - for layer_name in kv_cache_group.layer_names: - slot_mappings_by_layer[layer_name] = slot_mapping_gpu - - return attn_metadata, positions_cpu.to(device=self.device), slot_mappings_by_layer +# =================================================================== +# Code Predictor Transformer Model +# =================================================================== class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module): + """Transformer model for the code predictor (re-prefill, no KV cache).""" + def __init__( self, config: Qwen3TTSTalkerCodePredictorConfig, *, talker_hidden_size: int | None = None, - cache_config=None, - quant_config=None, prefix: str = "", ) -> None: super().__init__() self.config = config - self.quant_config = quant_config self.layers = nn.ModuleList( - [ - Qwen3DecoderLayer( - config, cache_config=cache_config, quant_config=quant_config, prefix=f"{prefix}.layers.{i}" - ) - for i in range(config.num_hidden_layers) - ] + [_CodePredictorDecoderLayer(config, prefix=f"{prefix}.layers.{i}") for i in range(config.num_hidden_layers)] ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - # Official code_predictor uses one embedding table per residual group. - # Some Qwen3-TTS checkpoints store codec embeddings in the talker hidden - # space, even when `code_predictor_config.hidden_size` is smaller. - # We keep the embedding dim aligned with the checkpoint and project down - # via `small_to_mtp_projection` in the wrapper module. + # Codec embeddings: one per residual group. Stored in talker hidden dim + # (some checkpoints use talker_hidden_size != code_predictor hidden_size). emb_dim = int(talker_hidden_size) if talker_hidden_size is not None else int(config.hidden_size) self.codec_embedding = nn.ModuleList( [nn.Embedding(config.vocab_size, emb_dim) for _ in range(config.num_code_groups - 1)] @@ -238,18 +220,18 @@ class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module): def get_input_embeddings(self) -> nn.ModuleList: return self.codec_embedding - def forward(self, positions: torch.Tensor, inputs_embeds: torch.Tensor) -> torch.Tensor: - # Token-major: [num_tokens, hidden] + def forward( + self, + inputs_embeds: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: hidden_states = inputs_embeds - residual = None for layer in self.layers: - hidden_states, residual = layer(positions, hidden_states, residual) - hidden_states, _ = self.norm(hidden_states, residual) + hidden_states = layer(hidden_states, position_ids) + hidden_states = self.norm(hidden_states) return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Match vLLM Qwen2/Qwen3 packing conventions: q_proj/k_proj/v_proj -> qkv_proj, - # gate_proj/up_proj -> gate_up_proj. stacked_params_mapping = [ ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), @@ -262,13 +244,6 @@ class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module): for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue - if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): - param = params_dict[scale_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0] - weight_loader(param, loaded_weight) - loaded_params.add(scale_name) - continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: @@ -309,8 +284,31 @@ class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module): return loaded_params +# =================================================================== +# Code Predictor Wrapper (optimized re-prefill + torch.compile) +# =================================================================== + + class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): - """vLLM-native code_predictor used by the AR talker (residual codebooks).""" + """vLLM-native code_predictor for the AR talker (residual codebooks). + + Re-prefill approach: each AR step forwards the full growing sequence + through the 5-layer transformer. No KV cache needed. This trades + ~O(T^2) extra attention FLOPs (negligible for T=16, 5 layers) for + zero KV cache management overhead and a simpler execution model. + + Optimizations over baseline: + 1. torch.compile on model forward -- fuses 60+ small kernel launches per step + into fewer fused kernels (4x speedup on model_fwd, ~75% of total time). + 2. Pre-allocated embedding buffer [B, max_seq, H] -- no torch.cat per step. + 3. Projection caching -- each token projected once and cached, avoids O(T^2) + redundant projections. + 4. Pre-allocated position_ids [max_seq] -- no torch.arange per step. + 5. Inline sampling -- no custom op / forward_context overhead. + 6. No context managers in forward(). + 7. Cached module references -- bypass nn.Module.__call__ and ModuleList indexing. + 8. Pre-allocated output tensor. + """ def __init__( self, @@ -325,16 +323,12 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): self.config = config self.talker_config = talker_config - # Keep module/weight names aligned with official checkpoint (talker.code_predictor.model.*). self.model = Qwen3TTSTalkerCodePredictorModelVLLM( config, talker_hidden_size=int(talker_config.hidden_size), - cache_config=vllm_config.cache_config, - quant_config=vllm_config.quant_config, prefix=f"{prefix}.model", ) - # One head per residual group. self.lm_head = nn.ModuleList( [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_code_groups - 1)] ) @@ -344,14 +338,21 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): else: self.small_to_mtp_projection = nn.Identity() - self._kv_cache: _LocalPredictorKVCache | None = None + self._num_groups = int(config.num_code_groups) + self._talker_hidden = int(talker_config.hidden_size) + self._cp_hidden = int(config.hidden_size) + + # Pre-allocated buffers (lazily initialized on first forward). + self._proj_buf: torch.Tensor | None = None + self._pos_ids: torch.Tensor | None = None + + # torch.compile: fuse small kernels in the 5-layer transformer. + self._compiled_model_fwd: object | None = None def get_input_embeddings(self) -> nn.ModuleList: return self.model.get_input_embeddings() def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Ensure all vLLM custom layers consult the predictor vllm_config - # (esp. for Attention static_forward_context). with set_current_vllm_config(self._vllm_config): loaded: set[str] = set() model_weights: list[tuple[str, torch.Tensor]] = [] @@ -373,95 +374,52 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): loaded.add(name) return loaded - def _maybe_init_kv_cache(self, device: torch.device) -> None: - if self._kv_cache is not None: + # ------------------------------------------------------------------ + # Pre-allocated buffer management + # ------------------------------------------------------------------ + + def _ensure_buffers(self, bsz: int, device: torch.device, dtype: torch.dtype) -> None: + max_seq = self._num_groups + 1 + if ( + self._proj_buf is not None + and self._proj_buf.shape[0] >= bsz + and self._proj_buf.device == device + and self._proj_buf.dtype == dtype + ): return - with set_current_vllm_config(self._vllm_config): - max_seq_len = int(getattr(self.config, "num_code_groups", 16) or 16) - # Upper bound on batch size: vLLM scheduler max_num_seqs (fallback 8). - max_batch = int(getattr(self._vllm_config.scheduler_config, "max_num_seqs", 8) or 8) - max_batch = max(1, max_batch) - self._kv_cache = _LocalPredictorKVCache( - vllm_config=self._vllm_config, - max_seq_len=max_seq_len, - max_batch_size=max_batch, - device=device, - ) - - @torch.inference_mode() - def reset_cache(self) -> None: - # We reuse a fixed kv cache buffer and overwrite starting at slot 0. - # No action required here (seq_lens controls what is read). - return - - @torch.inference_mode() - def prefill_logits(self, inputs_embeds: torch.Tensor) -> torch.Tensor: - """Prefill with 2 tokens: [past_hidden, layer0_embed]. Returns logits for residual group 0.""" - self._maybe_init_kv_cache(inputs_embeds.device) - assert self._kv_cache is not None - - bsz = int(inputs_embeds.shape[0]) - qlen = 2 - # Flatten to token-major. - model_dtype = next(self.model.parameters()).dtype - hs = inputs_embeds.to(dtype=model_dtype).reshape(bsz * qlen, -1) - hs = self.small_to_mtp_projection(hs) - - query_lens = torch.full((bsz,), qlen, dtype=torch.int32) - seq_lens = query_lens.clone() - attn_metadata, positions, slot_mappings = self._kv_cache.build_attn_metadata( - num_reqs=bsz, query_lens=query_lens, seq_lens=seq_lens + self._proj_buf = torch.zeros( + bsz, + max_seq, + self._cp_hidden, + dtype=dtype, + device=device, + ) + self._pos_ids = torch.arange( + max_seq, + dtype=torch.long, + device=device, ) - with ( - set_current_vllm_config(self._vllm_config), - set_forward_context( - attn_metadata, - self._vllm_config, - num_tokens=int(hs.shape[0]), - slot_mapping=slot_mappings, - ), - ): - out = self.model(positions=positions, inputs_embeds=hs) + def _setup_compile(self) -> None: + """Lazily set up torch.compiled model forward for kernel fusion. - # Gather last token per request. - last_idx = torch.arange(qlen - 1, bsz * qlen, step=qlen, device=out.device, dtype=torch.long) - last_h = out.index_select(0, last_idx) - logits = self.lm_head[0](last_h) - return logits - - @torch.inference_mode() - def decode_logits(self, input_ids: torch.Tensor, *, generation_step: int, past_seq_len: int) -> torch.Tensor: - """Decode one new token for residual group `generation_step` (1..Q-1).""" - self._maybe_init_kv_cache(input_ids.device) - assert self._kv_cache is not None - bsz = int(input_ids.shape[0]) - if generation_step <= 0: - raise ValueError("generation_step must be >= 1 for decode_logits") - - embed_idx = generation_step - 1 - hs = self.model.get_input_embeddings()[embed_idx](input_ids.to(dtype=torch.long).reshape(bsz, 1)) - hs = self.small_to_mtp_projection(hs.reshape(bsz, -1)) - - query_lens = torch.ones((bsz,), dtype=torch.int32) - seq_lens = torch.full((bsz,), int(past_seq_len) + 1, dtype=torch.int32) - attn_metadata, positions, slot_mappings = self._kv_cache.build_attn_metadata( - num_reqs=bsz, query_lens=query_lens, seq_lens=seq_lens + Uses ``mode="default"`` so Inductor performs operator fusion without + capturing its own CUDA graphs. This avoids conflicts with vLLM's + ``CUDAGraphWrapper`` which manages CUDA graphs for the main Talker + model on the default stream. + """ + if self._compiled_model_fwd is not None: + return + self._compiled_model_fwd = torch.compile( + self.model.forward, + mode="default", + dynamic=True, ) + logger.info("code_predictor: torch.compile enabled (mode=default)") - with ( - set_current_vllm_config(self._vllm_config), - set_forward_context( - attn_metadata, - self._vllm_config, - num_tokens=int(hs.shape[0]), - slot_mapping=slot_mappings, - ), - ): - out = self.model(positions=positions, inputs_embeds=hs) - - logits = self.lm_head[generation_step](out) - return logits + # ------------------------------------------------------------------ + # Optimized forward: re-prefill + torch.compile + projection cache + # ------------------------------------------------------------------ @torch.inference_mode() def forward( @@ -470,62 +428,71 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): layer0_embed: torch.Tensor, last_talker_hidden: torch.Tensor, do_sample: bool = True, - temperature: float = 1.0, + temperature: float = 0.9, top_k: int = 50, top_p: float = 1.0, ) -> torch.Tensor: - """Full autoregressive prediction of residual codebooks 1..Q-1. + """Predict residual codebooks 1..Q-1 autoregressively via re-prefill. - Args: - layer0_code: [B, 1] first-layer codec token ids. - layer0_embed: [B, 1, H] embedding of layer0_code. - last_talker_hidden: [B, 1, H] hidden state from the talker. - do_sample: whether to sample or take argmax. - temperature: sampling temperature. - top_k: top-k filtering. - top_p: top-p (nucleus) filtering. + torch.compile fuses the ~60 small kernel launches per step into fewer + fused kernels, reducing kernel launch overhead by ~75%. - Returns: - audio_codes: [B, Q] all codebook tokens (layer0 + residuals). + Projection caching: each token is projected once via small_to_mtp_projection + and cached in _proj_buf, avoiding redundant re-projection of past tokens. """ bsz = int(layer0_code.shape[0]) - num_groups = int(self.config.num_code_groups) - max_steps = num_groups - 1 + num_groups = self._num_groups + device = layer0_code.device + dtype = layer0_embed.dtype - # Reset KV cache for a fresh sequence. - self.reset_cache() + all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device) + all_codes[:, 0] = layer0_code.reshape(bsz) - # Prefill: feed [last_talker_hidden, layer0_embed] → logits for group 1. - prefill_input = torch.cat([last_talker_hidden, layer0_embed], dim=1) # [B, 2, H] - logits = self.prefill_logits(prefill_input) # [B, vocab] + self._ensure_buffers(bsz, device, dtype) + self._setup_compile() - all_codes = [layer0_code.reshape(bsz, 1)] - past_seq_len = 2 + proj_buf = self._proj_buf + pos_ids = self._pos_ids + + projection = self.small_to_mtp_projection + model_fwd = self._compiled_model_fwd + lm_heads = list(self.lm_head) + codec_embeds = list(self.model.codec_embedding) + + proj_buf[:bsz, 0, :] = projection(last_talker_hidden.reshape(bsz, 1, -1)).reshape(bsz, -1) + proj_buf[:bsz, 1, :] = projection(layer0_embed.reshape(bsz, 1, -1)).reshape(bsz, -1) + + use_sampling = do_sample and temperature > 0 + inv_temperature = 1.0 / max(temperature, 1e-6) if use_sampling else 0.0 + if use_sampling and top_p != 1.0: + raise NotImplementedError( + "top_p sampling is not implemented for the vLLM-native code predictor; please set top_p=1.0." + ) for step in range(1, num_groups): - # Sample or argmax from logits. - if do_sample and temperature > 0: - scaled = logits.float() / temperature + seq_len = step + 1 + + projected = proj_buf[:bsz, :seq_len, :] + step_pos_ids = pos_ids[:seq_len] if bsz == 1 else pos_ids[:seq_len].repeat(bsz) + + hidden_out = model_fwd(projected, step_pos_ids) + + logits = lm_heads[step - 1](hidden_out[:, -1, :]) + + if use_sampling: + scaled = logits * inv_temperature if top_k > 0: topk_vals, _ = scaled.topk(top_k, dim=-1) scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) - probs = torch.softmax(scaled, dim=-1) - # Guard against NaN/Inf from lower-precision dtypes. - probs = probs.clamp(min=0.0) - row_sums = probs.sum(dim=-1, keepdim=True) - probs = torch.where(row_sums > 0, probs / row_sums, torch.ones_like(probs) / probs.shape[-1]) - next_ids = torch.multinomial(probs, num_samples=1) # [B, 1] + probs = F.softmax(scaled, dim=-1) + next_ids = torch.multinomial(probs, num_samples=1) else: - next_ids = logits.argmax(dim=-1, keepdim=True) # [B, 1] - all_codes.append(next_ids) + next_ids = logits.argmax(dim=-1, keepdim=True) - # If not the last step, decode one more token. - if step < max_steps: - logits = self.decode_logits( - next_ids.reshape(bsz), - generation_step=step, - past_seq_len=past_seq_len, - ) - past_seq_len += 1 + all_codes[:, step] = next_ids.reshape(bsz) - return torch.cat(all_codes, dim=1) # [B, Q] + if step < num_groups - 1: + new_embed = codec_embeds[step - 1](next_ids) + proj_buf[:bsz, step + 1, :] = projection(new_embed.reshape(bsz, 1, -1)).reshape(bsz, -1) + + return all_codes From 32fd084ae25afa68a8f5e53dc79502a631674839 Mon Sep 17 00:00:00 2001 From: Baoyuan Qi Date: Fri, 6 Mar 2026 14:45:08 +0800 Subject: [PATCH 3/9] [MiMo-Audio] Bugfix tp lg than 1 (#1688) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 齐保元 --- vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py | 3 +++ vllm_omni/model_executor/stage_configs/mimo_audio.yaml | 2 ++ .../model_executor/stage_configs/mimo_audio_async_chunk.yaml | 2 ++ 3 files changed, 7 insertions(+) diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index 3b9f94f23..c1b71e076 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -556,12 +556,14 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor config.hidden_size, bias=False, return_bias=False, + gather_output=True, ) self.hidden_states_downcast = ColumnParallelLinear( config.hidden_size, self.local_config.hidden_size, bias=False, return_bias=False, + gather_output=True, ) self.lm_head = ColumnParallelLinear( @@ -569,6 +571,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor config.vocab_size, bias=False, return_bias=False, + gather_output=True, ) # Re-encode the sum of multi-layer RVQ embeddings to obtain true Audio Code Embeddings diff --git a/vllm_omni/model_executor/stage_configs/mimo_audio.yaml b/vllm_omni/model_executor/stage_configs/mimo_audio.yaml index 552ba2ab5..b824fcb41 100644 --- a/vllm_omni/model_executor/stage_configs/mimo_audio.yaml +++ b/vllm_omni/model_executor/stage_configs/mimo_audio.yaml @@ -15,6 +15,7 @@ stage_args: model_arch: MiMoAudioForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs) gpu_memory_utilization: 0.3 enforce_eager: true # need to discuss trust_remote_code: true @@ -45,6 +46,7 @@ stage_args: model_arch: MiMoAudioForConditionalGeneration worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs) gpu_memory_utilization: 0.2 enforce_eager: true trust_remote_code: true diff --git a/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml index a60fc566a..7177aa809 100644 --- a/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml +++ b/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml @@ -15,6 +15,7 @@ stage_args: model_arch: MiMoAudioForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs) gpu_memory_utilization: 0.3 enforce_eager: true # need to discuss trust_remote_code: true @@ -46,6 +47,7 @@ stage_args: model_arch: MiMoAudioForConditionalGeneration worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs) gpu_memory_utilization: 0.2 enforce_eager: true trust_remote_code: true From 3bc89de0c0f7ae0247977244d79e8905d80a65c9 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Fri, 6 Mar 2026 01:45:32 -0500 Subject: [PATCH 4/9] Add non-async chunk support for Qwen3-TTS (#1678) Signed-off-by: linyueqian --- .buildkite/test-nightly.yml | 21 ++ .../qwen3-tts/vllm_omni/bench_async_chunk.py | 301 ++++++++++++++++++ .../qwen3-tts/vllm_omni/plot_async_chunk.py | 249 +++++++++++++++ .../vllm_omni/run_async_chunk_benchmark.sh | 167 ++++++++++ tests/e2e/online_serving/test_qwen3_tts.py | 66 +++- .../qwen3_tts_no_async_chunk.yaml | 67 ++++ .../stage_input_processors/qwen3_tts.py | 31 ++ 7 files changed, 898 insertions(+), 4 deletions(-) create mode 100644 benchmarks/qwen3-tts/vllm_omni/bench_async_chunk.py create mode 100644 benchmarks/qwen3-tts/vllm_omni/plot_async_chunk.py create mode 100755 benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh create mode 100644 vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 8b21acfa7..8f141576d 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -62,6 +62,27 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test" + timeout_in_minutes: 30 + depends_on: upload-nightly-pipeline + if: build.env("NIGHTLY") == "1" + commands: + - | + export VLLM_WORKER_MULTIPROC_METHOD=spawn + pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk + agents: + queue: "gpu_4_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: ":full_moon: Omni Model Perf Test with H100" key: nightly-performance timeout_in_minutes: 180 diff --git a/benchmarks/qwen3-tts/vllm_omni/bench_async_chunk.py b/benchmarks/qwen3-tts/vllm_omni/bench_async_chunk.py new file mode 100644 index 000000000..3497ae821 --- /dev/null +++ b/benchmarks/qwen3-tts/vllm_omni/bench_async_chunk.py @@ -0,0 +1,301 @@ +"""Benchmark comparing async_chunk on vs off for Qwen3-TTS. + +Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across +concurrency levels for both async_chunk modes. Saves results as JSON. + +Usage: + # Run against a server already serving with a given config: + python bench_async_chunk.py \ + --host 127.0.0.1 --port 8000 \ + --config-name async_chunk_on \ + --num-prompts 50 \ + --max-concurrency 1 10 \ + --result-dir results/ +""" + +import argparse +import asyncio +import json +import time +from dataclasses import asdict, dataclass, field +from datetime import datetime +from pathlib import Path + +import aiohttp +import numpy as np +from tqdm.asyncio import tqdm + +PROMPTS = [ + "Hello, welcome to the voice synthesis benchmark test.", + "She said she would be here by noon, but nobody showed up.", + "The quick brown fox jumps over the lazy dog near the riverbank.", + "I can't believe how beautiful the sunset looks from up here on the mountain.", + "Please remember to bring your identification documents to the appointment tomorrow morning.", + "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?", + "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.", + "After the meeting, we should discuss the quarterly results and plan for the next phase.", + "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.", + "The train leaves at half past seven, so we need to arrive at the station before then.", + "Could you please turn down the music a little bit, I'm trying to concentrate on my work.", + "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.", +] + + +@dataclass +class RequestResult: + success: bool = False + ttfp: float = 0.0 + e2e: float = 0.0 + audio_bytes: int = 0 + audio_duration: float = 0.0 + rtf: float = 0.0 + prompt: str = "" + error: str = "" + + +@dataclass +class BenchmarkResult: + config_name: str = "" + concurrency: int = 0 + num_prompts: int = 0 + completed: int = 0 + failed: int = 0 + duration_s: float = 0.0 + mean_ttfp_ms: float = 0.0 + median_ttfp_ms: float = 0.0 + std_ttfp_ms: float = 0.0 + p90_ttfp_ms: float = 0.0 + p95_ttfp_ms: float = 0.0 + p99_ttfp_ms: float = 0.0 + mean_e2e_ms: float = 0.0 + median_e2e_ms: float = 0.0 + std_e2e_ms: float = 0.0 + p90_e2e_ms: float = 0.0 + p95_e2e_ms: float = 0.0 + p99_e2e_ms: float = 0.0 + mean_rtf: float = 0.0 + median_rtf: float = 0.0 + std_rtf: float = 0.0 + mean_audio_duration_s: float = 0.0 + total_audio_duration_s: float = 0.0 + audio_throughput: float = 0.0 + request_throughput: float = 0.0 + per_request: list = field(default_factory=list) + + +def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float: + return num_bytes / sample_width / sample_rate + + +async def send_tts_request( + session: aiohttp.ClientSession, + api_url: str, + prompt: str, + voice: str = "vivian", + language: str = "English", + stream: bool = True, + pbar: tqdm | None = None, +) -> RequestResult: + payload = { + "input": prompt, + "voice": voice, + "language": language, + "stream": stream, + "response_format": "pcm", + } + + result = RequestResult(prompt=prompt) + st = time.perf_counter() + + try: + async with session.post(api_url, json=payload) as response: + if response.status != 200: + result.error = f"HTTP {response.status}: {await response.text()}" + return result + + first_chunk = True + total_bytes = 0 + + async for chunk in response.content.iter_any(): + if first_chunk and len(chunk) > 0: + result.ttfp = time.perf_counter() - st + first_chunk = False + total_bytes += len(chunk) + + result.e2e = time.perf_counter() - st + result.audio_bytes = total_bytes + result.audio_duration = pcm_bytes_to_duration(total_bytes) + if result.audio_duration > 0: + result.rtf = result.e2e / result.audio_duration + result.success = True + + except Exception as e: + result.error = str(e) + result.e2e = time.perf_counter() - st + + if pbar: + pbar.update(1) + return result + + +async def run_benchmark( + host: str, + port: int, + num_prompts: int, + max_concurrency: int, + num_warmups: int = 3, + voice: str = "vivian", + language: str = "English", + stream: bool = True, +) -> BenchmarkResult: + api_url = f"http://{host}:{port}/v1/audio/speech" + + connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60) + session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600)) + + if num_warmups > 0: + print(f" Warming up with {num_warmups} requests...") + warmup_tasks = [ + send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream) + for i in range(num_warmups) + ] + await asyncio.gather(*warmup_tasks) + print(" Warmup done.") + + request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)] + + print(f" Running {num_prompts} requests with concurrency={max_concurrency}...") + semaphore = asyncio.Semaphore(max_concurrency) + pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}") + + async def limited_request(prompt): + async with semaphore: + return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar) + + start_time = time.perf_counter() + tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts] + results: list[RequestResult] = await asyncio.gather(*tasks) + duration = time.perf_counter() - start_time + pbar.close() + + await session.close() + + successful = [r for r in results if r.success] + failed = [r for r in results if not r.success] + + bench = BenchmarkResult( + concurrency=max_concurrency, + num_prompts=num_prompts, + completed=len(successful), + failed=len(failed), + duration_s=duration, + ) + + if successful: + ttfps = [r.ttfp * 1000 for r in successful] + e2es = [r.e2e * 1000 for r in successful] + rtfs = [r.rtf for r in successful] + audio_durs = [r.audio_duration for r in successful] + + bench.mean_ttfp_ms = float(np.mean(ttfps)) + bench.median_ttfp_ms = float(np.median(ttfps)) + bench.std_ttfp_ms = float(np.std(ttfps)) + bench.p90_ttfp_ms = float(np.percentile(ttfps, 90)) + bench.p95_ttfp_ms = float(np.percentile(ttfps, 95)) + bench.p99_ttfp_ms = float(np.percentile(ttfps, 99)) + + bench.mean_e2e_ms = float(np.mean(e2es)) + bench.median_e2e_ms = float(np.median(e2es)) + bench.std_e2e_ms = float(np.std(e2es)) + bench.p90_e2e_ms = float(np.percentile(e2es, 90)) + bench.p95_e2e_ms = float(np.percentile(e2es, 95)) + bench.p99_e2e_ms = float(np.percentile(e2es, 99)) + + bench.mean_rtf = float(np.mean(rtfs)) + bench.median_rtf = float(np.median(rtfs)) + bench.std_rtf = float(np.std(rtfs)) + + bench.mean_audio_duration_s = float(np.mean(audio_durs)) + bench.total_audio_duration_s = float(np.sum(audio_durs)) + bench.audio_throughput = bench.total_audio_duration_s / duration + bench.request_throughput = len(successful) / duration + + bench.per_request = [ + { + "ttfp_ms": r.ttfp * 1000, + "e2e_ms": r.e2e * 1000, + "rtf": r.rtf, + "audio_duration_s": r.audio_duration, + "prompt": r.prompt, + } + for r in successful + ] + + print(f"\n{'=' * 60}") + print(f" Concurrency: {max_concurrency} | Completed: {bench.completed} | Failed: {bench.failed}") + print(f" Duration: {duration:.2f}s | Throughput: {bench.request_throughput:.2f} req/s") + print( + f" TTFP (ms): mean={bench.mean_ttfp_ms:.1f} median={bench.median_ttfp_ms:.1f}" + f" p90={bench.p90_ttfp_ms:.1f} p99={bench.p99_ttfp_ms:.1f}" + ) + print( + f" E2E (ms): mean={bench.mean_e2e_ms:.1f} median={bench.median_e2e_ms:.1f}" + f" p90={bench.p90_e2e_ms:.1f} p99={bench.p99_e2e_ms:.1f}" + ) + print(f" RTF: mean={bench.mean_rtf:.3f} median={bench.median_rtf:.3f}") + print(f" Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec") + print(f"{'=' * 60}\n") + + if failed: + for r in failed[:3]: + print(f" [ERROR] {r.error[:200]}") + + return bench + + +async def main(args): + all_results = [] + + for concurrency in args.max_concurrency: + result = await run_benchmark( + host=args.host, + port=args.port, + num_prompts=args.num_prompts, + max_concurrency=concurrency, + num_warmups=args.num_warmups, + voice=args.voice, + language=args.language, + stream=args.stream, + ) + result.config_name = args.config_name + all_results.append(asdict(result)) + + result_dir = Path(args.result_dir) + result_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json" + + with open(result_file, "w") as f: + json.dump(all_results, f, indent=2) + print(f"Results saved to {result_file}") + + +def parse_args(): + parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client") + parser.add_argument("--host", type=str, default="127.0.0.1") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--num-prompts", type=int, default=50) + parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10]) + parser.add_argument("--num-warmups", type=int, default=3) + parser.add_argument("--voice", type=str, default="vivian") + parser.add_argument("--language", type=str, default="English") + parser.add_argument("--stream", action="store_true", default=True) + parser.add_argument("--no-stream", dest="stream", action="store_false") + parser.add_argument("--config-name", type=str, default="async_chunk_on") + parser.add_argument("--result-dir", type=str, default="results") + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + asyncio.run(main(args)) diff --git a/benchmarks/qwen3-tts/vllm_omni/plot_async_chunk.py b/benchmarks/qwen3-tts/vllm_omni/plot_async_chunk.py new file mode 100644 index 000000000..dd03d9626 --- /dev/null +++ b/benchmarks/qwen3-tts/vllm_omni/plot_async_chunk.py @@ -0,0 +1,249 @@ +"""Plot TTFP comparison: async_chunk off vs on. + +Generates a bar chart with improvement arrows, matching the Qwen3-Omni +async_chunk benchmark figure style. + +Usage: + python plot_async_chunk.py \ + --off results/bench_async_chunk_off_*.json \ + --on results/bench_async_chunk_on_*.json \ + --output results/qwen3_tts_async_chunk_ttfp.png + + # Also supports E2E and RTF metrics: + python plot_async_chunk.py \ + --off results/bench_async_chunk_off_*.json \ + --on results/bench_async_chunk_on_*.json \ + --metric e2e \ + --output results/qwen3_tts_async_chunk_e2e.png +""" + +import argparse +import json +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +METRIC_CONFIG = { + "ttfp": { + "key": "mean_ttfp_ms", + "ylabel": "TTFP (s)", + "title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency", + "to_seconds": True, + }, + "e2e": { + "key": "mean_e2e_ms", + "ylabel": "E2E (s)", + "title": "E2E Latency - Qwen3-TTS, by concurrency", + "to_seconds": True, + }, + "rtf": { + "key": "mean_rtf", + "ylabel": "RTF", + "title": "Real-Time Factor - Qwen3-TTS, by concurrency", + "to_seconds": False, + }, +} + + +def load_results(path: str) -> list[dict]: + with open(path) as f: + return json.load(f) + + +def plot_ttfp_comparison( + off_results: list[dict], + on_results: list[dict], + metric: str, + output_path: str, + title_override: str | None = None, +): + cfg = METRIC_CONFIG[metric] + key = cfg["key"] + to_seconds = cfg["to_seconds"] + + off_map = {r["concurrency"]: r for r in off_results} + on_map = {r["concurrency"]: r for r in on_results} + concurrencies = sorted(set(off_map.keys()) & set(on_map.keys())) + + off_vals = [] + on_vals = [] + for c in concurrencies: + v_off = off_map[c][key] + v_on = on_map[c][key] + if to_seconds: + v_off /= 1000.0 + v_on /= 1000.0 + off_vals.append(v_off) + on_vals.append(v_on) + + fig, ax = plt.subplots(figsize=(8, 6)) + + x = np.arange(len(concurrencies)) + width = 0.3 + + ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none") + ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8") + + # Draw improvement arrows and labels + for i in range(len(concurrencies)): + v_off = off_vals[i] + v_on = on_vals[i] + if v_on > 0: + improvement = v_off / v_on + else: + improvement = float("inf") + + # Arrow from top of off-bar to top of on-bar + arrow_start_x = x[i] - width / 2 + arrow_start_y = v_off * 0.95 + arrow_end_x = x[i] + width / 2 + arrow_end_y = v_on * 1.05 + + ax.annotate( + "", + xy=(arrow_end_x, arrow_end_y), + xytext=(arrow_start_x, arrow_start_y), + arrowprops=dict(arrowstyle="->", color="red", lw=1.5), + ) + + # Improvement label + label_x = (arrow_start_x + arrow_end_x) / 2 + label_y = arrow_start_y + (v_off - v_on) * 0.15 + ax.text( + label_x, + label_y, + f"{improvement:.1f}x improvement", + ha="center", + va="bottom", + fontsize=10, + color="red", + fontweight="bold", + ) + + title = title_override or cfg["title"] + ax.set_title(title, fontsize=13, fontweight="bold") + ax.set_ylabel(cfg["ylabel"], fontsize=12) + ax.set_xlabel("Max concurrency", fontsize=12) + ax.set_xticks(x) + ax.set_xticklabels([str(c) for c in concurrencies]) + ax.set_yscale("log") + ax.legend(loc="upper left", fontsize=11) + ax.grid(axis="y", alpha=0.3, linestyle="--") + ax.set_axisbelow(True) + + plt.tight_layout() + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + plt.savefig(output_path, dpi=150, bbox_inches="tight") + print(f"Plot saved to {output_path}") + plt.close() + + +def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str): + """Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons.""" + off_map = {r["concurrency"]: r for r in off_results} + on_map = {r["concurrency"]: r for r in on_results} + concurrencies = sorted(set(off_map.keys()) & set(on_map.keys())) + + fig, axes = plt.subplots(1, 3, figsize=(18, 6)) + fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold") + + for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]): + cfg = METRIC_CONFIG[metric] + key = cfg["key"] + to_seconds = cfg["to_seconds"] + + off_vals = [] + on_vals = [] + for c in concurrencies: + v_off = off_map[c][key] + v_on = on_map[c][key] + if to_seconds: + v_off /= 1000.0 + v_on /= 1000.0 + off_vals.append(v_off) + on_vals.append(v_on) + + x = np.arange(len(concurrencies)) + width = 0.3 + ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB") + ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8") + + for i in range(len(concurrencies)): + if on_vals[i] > 0: + improvement = off_vals[i] / on_vals[i] + ax.annotate( + "", + xy=(x[i] + width / 2, on_vals[i] * 1.05), + xytext=(x[i] - width / 2, off_vals[i] * 0.95), + arrowprops=dict(arrowstyle="->", color="red", lw=1.5), + ) + label_y = off_vals[i] * 0.85 + ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold") + + ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold") + ax.set_ylabel(cfg["ylabel"], fontsize=11) + ax.set_xlabel("Max concurrency", fontsize=11) + ax.set_xticks(x) + ax.set_xticklabels([str(c) for c in concurrencies]) + if metric != "rtf": + ax.set_yscale("log") + ax.legend(fontsize=9) + ax.grid(axis="y", alpha=0.3, linestyle="--") + ax.set_axisbelow(True) + + plt.tight_layout() + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + plt.savefig(output_path, dpi=150, bbox_inches="tight") + print(f"Plot saved to {output_path}") + plt.close() + + +def print_table(off_results: list[dict], on_results: list[dict]): + off_map = {r["concurrency"]: r for r in off_results} + on_map = {r["concurrency"]: r for r in on_results} + concurrencies = sorted(set(off_map.keys()) & set(on_map.keys())) + + print("\n## Benchmark Results: async_chunk off vs on\n") + print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |") + print("| --- | --- | --- | --- | --- |") + + for name, key, fmt in [ + ("TTFP (ms)", "mean_ttfp_ms", ".1f"), + ("E2E (ms)", "mean_e2e_ms", ".1f"), + ("RTF", "mean_rtf", ".3f"), + ("Throughput", "audio_throughput", ".2f"), + ]: + for c in concurrencies: + v_off = off_map[c].get(key, 0) + v_on = on_map[c].get(key, 0) + if v_on > 0 and key != "audio_throughput": + ratio = f"{v_off / v_on:.1f}x" + elif v_off > 0 and key == "audio_throughput": + ratio = f"{v_on / v_off:.1f}x" + else: + ratio = "N/A" + print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |") + + +def parse_args(): + parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS") + parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off") + parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on") + parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"]) + parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png") + parser.add_argument("--title", type=str, default=None, help="Custom title override") + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + off_results = load_results(args.off) + on_results = load_results(args.on) + + print_table(off_results, on_results) + + if args.metric == "all": + plot_all_metrics(off_results, on_results, args.output) + else: + plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title) diff --git a/benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh b/benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh new file mode 100755 index 000000000..61cf7757a --- /dev/null +++ b/benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh @@ -0,0 +1,167 @@ +#!/bin/bash +# Qwen3-TTS async_chunk on vs off Benchmark +# +# Starts two servers (async_chunk on and off), benchmarks both, +# and generates comparison plots. +# +# Usage: +# bash run_async_chunk_benchmark.sh +# +# Environment variables: +# GPU_DEVICE - GPU index (default: 0) +# MODEL - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice) +# NUM_PROMPTS - Prompts per concurrency level (default: 50) +# CONCURRENCY - Space-separated concurrency levels (default: "1 10") +# PORT_ON - Port for async_chunk on server (default: 8000) +# PORT_OFF - Port for async_chunk off server (default: 8001) + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +cd "$PROJECT_ROOT" + +GPU_DEVICE="${GPU_DEVICE:-0}" +MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}" +NUM_PROMPTS="${NUM_PROMPTS:-50}" +CONCURRENCY="${CONCURRENCY:-1 10}" +NUM_WARMUPS="${NUM_WARMUPS:-3}" +PORT_ON="${PORT_ON:-8000}" +PORT_OFF="${PORT_OFF:-8001}" +RESULT_DIR="${SCRIPT_DIR}/results" +TIMESTAMP="$(date +%Y%m%d_%H%M%S)" + +STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml" +STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml" + +mkdir -p "${RESULT_DIR}" + +echo "============================================================" +echo " Qwen3-TTS async_chunk Benchmark" +echo "============================================================" +echo " GPU: ${GPU_DEVICE}" +echo " Model: ${MODEL}" +echo " Prompts: ${NUM_PROMPTS}" +echo " Concurrency: ${CONCURRENCY}" +echo " Port (on/off): ${PORT_ON} / ${PORT_OFF}" +echo " Results: ${RESULT_DIR}" +echo "============================================================" + +cleanup() { + echo "Cleaning up servers..." + kill "$PID_ON" 2>/dev/null || true + kill "$PID_OFF" 2>/dev/null || true + wait "$PID_ON" 2>/dev/null || true + wait "$PID_OFF" 2>/dev/null || true +} +trap cleanup EXIT + +wait_for_server() { + local port=$1 + local name=$2 + local max_wait=300 + local elapsed=0 + echo "Waiting for ${name} server on port ${port}..." + while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do + sleep 5 + elapsed=$((elapsed + 5)) + if [ $elapsed -ge $max_wait ]; then + echo "ERROR: ${name} server failed to start within ${max_wait}s" + exit 1 + fi + done + echo "${name} server ready (${elapsed}s)" +} + +# ---- Phase 1: Start async_chunk ON server ---- +echo "" +echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..." +CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \ + --stage-configs-path "${STAGE_CONFIG_ON}" \ + --host 0.0.0.0 --port "${PORT_ON}" \ + --trust-remote-code --enforce-eager --omni \ + > "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 & +PID_ON=$! + +wait_for_server "${PORT_ON}" "async_chunk_on" + +echo "[Phase 1] Benchmarking async_chunk ON..." +# shellcheck disable=SC2086 +python "${SCRIPT_DIR}/bench_async_chunk.py" \ + --host 127.0.0.1 --port "${PORT_ON}" \ + --config-name "async_chunk_on" \ + --num-prompts "${NUM_PROMPTS}" \ + --max-concurrency ${CONCURRENCY} \ + --num-warmups "${NUM_WARMUPS}" \ + --result-dir "${RESULT_DIR}" + +echo "[Phase 1] Stopping async_chunk ON server..." +kill "$PID_ON" 2>/dev/null || true +wait "$PID_ON" 2>/dev/null || true +sleep 5 + +# ---- Phase 2: Start async_chunk OFF server ---- +echo "" +echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..." +CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \ + --stage-configs-path "${STAGE_CONFIG_OFF}" \ + --host 0.0.0.0 --port "${PORT_OFF}" \ + --trust-remote-code --enforce-eager --omni \ + > "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 & +PID_OFF=$! + +wait_for_server "${PORT_OFF}" "async_chunk_off" + +echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..." +# shellcheck disable=SC2086 +python "${SCRIPT_DIR}/bench_async_chunk.py" \ + --host 127.0.0.1 --port "${PORT_OFF}" \ + --config-name "async_chunk_off" \ + --num-prompts "${NUM_PROMPTS}" \ + --max-concurrency ${CONCURRENCY} \ + --num-warmups "${NUM_WARMUPS}" \ + --no-stream \ + --result-dir "${RESULT_DIR}" + +echo "[Phase 2] Stopping async_chunk OFF server..." +kill "$PID_OFF" 2>/dev/null || true +wait "$PID_OFF" 2>/dev/null || true + +# ---- Phase 3: Plot results ---- +echo "" +echo "[Phase 3] Generating plots..." + +# Find the latest result files +RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1) +RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1) + +if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then + echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/" + exit 1 +fi + +echo " ON results: ${RESULT_ON}" +echo " OFF results: ${RESULT_OFF}" + +# TTFP comparison (main figure) +python "${SCRIPT_DIR}/plot_async_chunk.py" \ + --off "${RESULT_OFF}" \ + --on "${RESULT_ON}" \ + --metric ttfp \ + --output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png" + +# All metrics comparison +python "${SCRIPT_DIR}/plot_async_chunk.py" \ + --off "${RESULT_OFF}" \ + --on "${RESULT_ON}" \ + --metric all \ + --output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png" + +echo "" +echo "============================================================" +echo " Benchmark complete!" +echo " Results: ${RESULT_DIR}/" +echo " Plots:" +echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png" +echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png" +echo "============================================================" diff --git a/tests/e2e/online_serving/test_qwen3_tts.py b/tests/e2e/online_serving/test_qwen3_tts.py index b4ca57835..2ab5cdd4b 100644 --- a/tests/e2e/online_serving/test_qwen3_tts.py +++ b/tests/e2e/online_serving/test_qwen3_tts.py @@ -23,11 +23,9 @@ from tests.utils import hardware_test MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" -def get_stage_config(): +def get_stage_config(name: str = "qwen3_tts.yaml"): """Get the stage config path for Qwen3-TTS.""" - return str( - Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "qwen3_tts.yaml" - ) + return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name) @pytest.fixture(scope="module") @@ -225,3 +223,63 @@ class TestQwen3TTSAPIEndpoints: data = response.json() assert "data" in data assert len(data["data"]) > 0 + + +@pytest.fixture(scope="module") +def omni_server_no_async_chunk(): + """Start vLLM-Omni server with non-async-chunk config.""" + stage_config_path = get_stage_config("qwen3_tts_no_async_chunk.yaml") + + with OmniServer( + MODEL, + [ + "--stage-configs-path", + stage_config_path, + "--stage-init-timeout", + "120", + "--trust-remote-code", + "--enforce-eager", + "--disable-log-stats", + ], + ) as server: + yield server + + +class TestQwen3TTSNoAsyncChunk: + """E2E tests for Qwen3-TTS in non-async-chunk (full decode) mode.""" + + @pytest.mark.core_model + @pytest.mark.omni + @hardware_test(res={"cuda": "L4"}, num_cards=4) + def test_speech_english(self, omni_server_no_async_chunk) -> None: + """Test English TTS with non-async-chunk pipeline.""" + response = make_speech_request( + host=omni_server_no_async_chunk.host, + port=omni_server_no_async_chunk.port, + text="Hello, how are you?", + voice="vivian", + language="English", + ) + + assert response.status_code == 200, f"Request failed: {response.text}" + assert response.headers.get("content-type") == "audio/wav" + assert verify_wav_audio(response.content), "Response is not valid WAV audio" + assert len(response.content) > MIN_AUDIO_BYTES + + @pytest.mark.core_model + @pytest.mark.omni + @hardware_test(res={"cuda": "L4"}, num_cards=4) + def test_speech_chinese(self, omni_server_no_async_chunk) -> None: + """Test Chinese TTS with non-async-chunk pipeline.""" + response = make_speech_request( + host=omni_server_no_async_chunk.host, + port=omni_server_no_async_chunk.port, + text="你好，我是通义千问", + voice="vivian", + language="Chinese", + ) + + assert response.status_code == 200, f"Request failed: {response.text}" + assert response.headers.get("content-type") == "audio/wav" + assert verify_wav_audio(response.content), "Response is not valid WAV audio" + assert len(response.content) > MIN_AUDIO_BYTES diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml new file mode 100644 index 000000000..da22953bb --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml @@ -0,0 +1,67 @@ +async_chunk: false +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: qwen3_tts + model_arch: Qwen3TTSTalkerForConditionalGeneration + hf_overrides: + architectures: [Qwen3TTSTalkerForConditionalGeneration] + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + enforce_eager: false + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: latent + gpu_memory_utilization: 0.3 + distributed_executor_backend: "mp" + max_num_batched_tokens: 512 + max_model_len: 4096 + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + seed: 42 + detokenize: false + repetition_penalty: 1.05 + stop_token_ids: [2150] + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: code2wav + model_arch: Qwen3TTSCode2Wav + hf_overrides: + architectures: [Qwen3TTSCode2Wav] + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio + gpu_memory_utilization: 0.2 + distributed_executor_backend: "mp" + max_num_batched_tokens: 8192 + max_model_len: 32768 + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav + final_output: true + final_output_type: audio + tts_args: + max_instructions_length: 500 + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + detokenize: true + repetition_penalty: 1.0 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index c0bcfa612..f91d246c6 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -8,6 +8,37 @@ from vllm.logger import init_logger logger = init_logger(__name__) +def talker2code2wav( + stage_list: list[Any], + engine_input_source: list[int], + prompt: Any = None, + requires_multimodal_data: bool = False, +) -> list[Any]: + """Non-async processor: wait for talker to finish, then pass all codes to code2wav at once.""" + from vllm_omni.inputs.data import OmniTokensPrompt + from vllm_omni.model_executor.stage_input_processors.qwen3_omni import _validate_stage_inputs + + talker_outputs = _validate_stage_inputs(stage_list, engine_input_source) + code2wav_inputs: list[OmniTokensPrompt] = [] + for talker_output in talker_outputs: + output = talker_output.outputs[0] + # audio_codes shape: [num_frames, Q] where Q=num_quantizers (16) + audio_codes = output.multimodal_output["audio_codes"].to(torch.long) + # Filter zero-padded frames (EOS/invalid steps), matching _extract_last_frame behavior + valid_mask = audio_codes.any(dim=1) + audio_codes = audio_codes[valid_mask] + # Code2Wav expects codebook-major flat: [Q*num_frames] + codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist() + code2wav_inputs.append( + OmniTokensPrompt( + prompt_token_ids=codec_codes, + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + return code2wav_inputs + + def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None: audio_codes = pooling_output.get("audio_codes") if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: From 846167b23766f6294970798bd18e7bbb8f14105c Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Fri, 6 Mar 2026 14:45:53 +0800 Subject: [PATCH 5/9] [1/N][Refactor] Clean up dead code in output processor (#1579) Signed-off-by: gcanlin --- vllm_omni/engine/output_processor.py | 158 ++---------------- .../entrypoints/openai/serving_speech.py | 2 +- 2 files changed, 19 insertions(+), 141 deletions(-) diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 697622f12..a37e84097 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -1,4 +1,3 @@ -from collections.abc import Callable from typing import Any import numpy as np @@ -60,9 +59,11 @@ class OmniRequestState(RequestState): incoming: dict[str, Any] = {} target_key = self.mm_type or "hidden" - # Iterate directly without unnecessary dict copy for k, v in payload.items(): - # Optional remap: if producer used "model_outputs" or "hidden", rename to mm_type + # Normalize producer keys to the modality name. + # AR runners produce {"hidden": ...} and generation + # runners produce {"model_outputs": ...}; remap both + # to the semantic modality key (e.g. "audio", "latent"). if k == "model_outputs": k = target_key elif k == "hidden" and target_key != "hidden": @@ -247,16 +248,17 @@ class OmniRequestState(RequestState): class MultimodalOutputProcessor(VLLMOutputProcessor): - """Handles multimodal output processing by normalizing EngineCoreOutput - before delegating to the base vLLM OutputProcessor. + """Handles multimodal output processing by capturing pooling_output + from EngineCoreOutput and accumulating it as multimodal tensors, + before delegating to the base vLLM OutputProcessor for text handling. - Strategy: - - Route by EngineCoreOutput.output_type when present - ("image", "text+image", "latents", "text"). - - Fallback to pooling/text heuristics when output_type is absent. - - Mutate EngineCoreOutput in-place to ensure vLLM's base processor can - produce the correct RequestOutput/PoolingRequestOutput. - - Allow custom per-modality handlers via register_handler(). + The actual data flow is: + 1. For each EngineCoreOutput with pooling_output and a detokenizer: + - Capture pooling_output into OmniRequestState.add_multimodal_tensor() + - Clear eco.pooling_output to force text path in base processor + 2. Base vLLM OutputProcessor handles text detokenization + 3. On finish, _consolidate_multimodal_tensors() concatenates accumulated tensors + 4. _new_completion_output() attaches mm_accumulated to CompletionOutput """ def __init__( @@ -273,27 +275,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): log_stats: Whether to log statistics stream_interval: Stream interval for output generation engine_core_output_type: Optional output type specification - (e.g., "image", "audio", "latents"). Used to route outputs - to appropriate processors. If None, output type is inferred. + (e.g., "image", "audio", "latent"). Used to tag multimodal + outputs with the correct modality key. """ super().__init__(tokenizer=tokenizer, log_stats=log_stats, stream_interval=stream_interval) - self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {} - self._reqid_to_mm_type: dict[str, str] = {} self.engine_core_output_type = engine_core_output_type - def register_handler(self, modality: str, handler: Callable[[EngineCoreOutput], None]) -> None: - """Register a custom handler for a specific modality. - - Allows custom processing logic for specific output modalities. - The handler is called before default processing for outputs - matching the specified modality. - - Args: - modality: Modality name (e.g., "image", "audio", "latents") - handler: Callable that takes an EngineCoreOutput and processes it - """ - self.output_handlers[modality.lower()] = handler - def add_request( self, request: EngineCoreRequest, @@ -346,20 +333,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): engine_core_timestamp: float | None = None, iteration_stats: IterationStats | None = None, ) -> OutputProcessorOutput: - self._reqid_to_mm_type.clear() for eco in engine_core_outputs: - mm_type = (self.engine_core_output_type or "").lower() - if mm_type: - self._reqid_to_mm_type[eco.request_id] = mm_type - self._route_and_normalize(eco) req_state = self.request_states.get(eco.request_id) if req_state is None or not isinstance(req_state, OmniRequestState): continue if eco.pooling_output is not None and req_state.detokenizer is not None: - req_state.add_multimodal_tensor( - eco.pooling_output, - (getattr(eco, "output_type", self.engine_core_output_type) or "").lower(), - ) + mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower() + req_state.add_multimodal_tensor(eco.pooling_output, mm_type) # Force text path in base processor for multimodal outputs. eco.pooling_output = None @@ -368,105 +348,3 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): engine_core_timestamp=engine_core_timestamp, iteration_stats=iteration_stats, ) - - # ---- routing helpers ---- - def _route_and_normalize(self, eco: EngineCoreOutput) -> None: - output_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower() - - # Custom handler first (if registered) - if output_type in self.output_handlers: - try: - self.output_handlers[output_type](eco) - # Fall through to default fixups in case the handler left gaps - except Exception: - logger.exception("Error in custom output handler for %s", output_type) - - if output_type == "image": - self._process_image_output(eco) - elif output_type in ("text+image", "text,image", "image+text"): - self._process_text_image_output(eco) - elif output_type in ("latents", "latent"): - self._process_latents_output(eco) - elif output_type in ("audio", "speech"): - self._process_audio_output(eco) - elif output_type == "text": - self._process_text_output(eco) - else: - # Fallback heuristic - if eco.pooling_output is not None: - self._process_pooling_output(eco) - else: - self._process_text_output(eco) - - # ---- modality processors ---- - def _process_image_output(self, eco: EngineCoreOutput) -> None: - """Ensure image tensors are surfaced via pooling_output for vLLM.""" - if eco.pooling_output is None: - tensor = self._extract_from_multimodal_outputs(eco, keys=("image", "images", "pixel_values", "pixels")) - if tensor is not None: - eco.pooling_output = tensor - - def _process_text_image_output(self, eco: EngineCoreOutput) -> None: - """Allow text+image outputs. Text path stays as new_token_ids; - image/latents route via pooling_output.""" - # Preserve text tokens as-is; ensure pooling_output carries image/latents - if eco.pooling_output is None: - tensor = self._extract_from_multimodal_outputs( - eco, - keys=( - "image", - "images", - "pixel_values", - "pixels", - "latent", - "latents", - "z", - ), - ) - if tensor is not None: - eco.pooling_output = tensor - - def _process_latents_output(self, eco: EngineCoreOutput) -> None: - """Ensure latent tensors are surfaced via pooling_output.""" - if eco.pooling_output is None: - tensor = self._extract_from_multimodal_outputs(eco, keys=("latent", "latents", "z", "posterior")) - if tensor is not None: - eco.pooling_output = tensor - - def _process_audio_output(self, eco: EngineCoreOutput) -> None: - """Ensure audio tensors are surfaced via pooling_output.""" - if eco.pooling_output is None: - tensor = self._extract_from_multimodal_outputs( - eco, keys=("audio", "audios", "wav", "waveform", "audio_pcm", "pcm") - ) - if tensor is not None: - eco.pooling_output = tensor - - def _process_text_output(self, eco: EngineCoreOutput) -> None: - """No-op; base processor will detokenize new_token_ids → text.""" - return - - def _process_pooling_output(self, eco: EngineCoreOutput) -> None: - """Optional sanity checks for pooling tensor.""" - if eco.pooling_output is None: - return - if not isinstance(eco.pooling_output, torch.Tensor): - # Best-effort: convert to tensor if it's a list/ndarray-like - try: - eco.pooling_output = torch.as_tensor(eco.pooling_output) - except Exception: - pass - - def _extract_from_multimodal_outputs(self, eco: EngineCoreOutput, keys: tuple[str, ...]) -> torch.Tensor | None: - mm = getattr(eco, "multimodal_outputs", None) - if not isinstance(mm, dict): - return None - for k in keys: - v = mm.get(k) - if isinstance(v, torch.Tensor): - return v - # Try the first tensor in the dict as a fallback - for v in mm.values(): - if isinstance(v, torch.Tensor): - return v - return None diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 9ff7bc767..8c292cc7a 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -418,7 +418,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): mm = getattr(ro, "multimodal_output", None) if ro else None if not mm: return None, None - key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None) + key = "audio" if "audio" in mm else None return mm, key def _build_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: From 92adcdeab45a1a2a41887411b208ca9aba50a734 Mon Sep 17 00:00:00 2001 From: JohnJan Date: Fri, 6 Mar 2026 14:46:57 +0800 Subject: [PATCH 6/9] [feature]: support flux2.klein cache_dit (#1209) Co-authored-by: Hongsheng Liu --- docs/user_guide/diffusion_acceleration.md | 1 + .../diffusion/cache/cache_dit_backend.py | 79 +++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 5b86562b3..462ab8b18 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -68,6 +68,7 @@ The following table shows which models are currently supported by each accelerat | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | | **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | +| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ### VideoGen diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 5008a8fa3..1c98d0724 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -346,6 +346,84 @@ def enable_cache_for_flux(pipeline: Any, cache_config: Any) -> Callable[[int], N return refresh_cache_context +def enable_cache_for_flux2_klein(pipeline: Any, cache_config: Any) -> Callable[[int], None]: + """Enable cache-dit for FLUX.2-klein-4B pipeline. + + Args: + pipeline: The FLUX.2-klein-4B pipeline instance. + cache_config: DiffusionCacheConfig instance with cache configuration. + Returns: + A refresh function that can be called with a new ``num_inference_steps`` + to update the cache context for the pipeline. + """ + # Build DBCacheConfig for transformer + db_cache_config = _build_db_cache_config(cache_config) + # The Fn_compute_blocks = 2 override is the most important decision here, + # and the rationale (quality degradation at Fn=1) only lives in flux2_klein. + db_cache_config.Fn_compute_blocks = 2 + + calibrator = None + if cache_config.enable_taylorseer: + taylorseer_order = cache_config.taylorseer_order + calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) + logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + + # Build ParamsModifier for transformer + modifier = ParamsModifier( + cache_config=db_cache_config, + calibrator_config=calibrator, + ) + + logger.info( + f"Enabling cache-dit on Flux2-Klein transformer with BlockAdapter: " + f"Fn={db_cache_config.Fn_compute_blocks}, " + f"Bn={db_cache_config.Bn_compute_blocks}, " + f"W={db_cache_config.max_warmup_steps}, " + ) + + # Enable cache-dit using BlockAdapter for transformer + cache_dit.enable_cache( + ( + BlockAdapter( + transformer=pipeline.transformer, + blocks=[ + pipeline.transformer.transformer_blocks, + pipeline.transformer.single_transformer_blocks, + ], + forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2], + params_modifiers=[modifier], + ) + ), + cache_config=db_cache_config, + ) + + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh cache context for the transformer with new num_inference_steps. + + Args: + pipeline: The FLUX.2-klein-4B pipeline instance. + num_inference_steps: New number of inference steps. + """ + if cache_config.scm_steps_mask_policy is None: + cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) + else: + cache_dit.refresh_context( + pipeline.transformer, + cache_config=DBCacheConfig().reset( + num_inference_steps=num_inference_steps, + steps_computation_mask=cache_dit.steps_mask( + mask_policy=cache_config.scm_steps_mask_policy, + total_steps=num_inference_steps, + ), + Fn_compute_blocks=db_cache_config.Fn_compute_blocks, + steps_computation_policy=cache_config.scm_steps_policy, + ), + verbose=verbose, + ) + + return refresh_cache_context + + def enable_cache_for_sd3(pipeline: Any, cache_config: Any) -> Callable[[int], None]: """Enable cache-dit for StableDiffusion3Pipeline. @@ -913,6 +991,7 @@ CUSTOM_DIT_ENABLERS.update( "Wan22I2VPipeline": enable_cache_for_wan22, "Wan22TI2VPipeline": enable_cache_for_wan22, "FluxPipeline": enable_cache_for_flux, + "Flux2KleinPipeline": enable_cache_for_flux2_klein, "LongCatImagePipeline": enable_cache_for_longcat_image, "LongCatImageEditPipeline": enable_cache_for_longcat_image, "StableDiffusion3Pipeline": enable_cache_for_sd3, From 411fe05bd47de2459e18f33ffcb561737cc75538 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Fri, 6 Mar 2026 01:49:19 -0500 Subject: [PATCH 7/9] [skip CI][Docs] Add TTS model developer guide (#1693) Signed-off-by: linyueqian --- docs/.nav.yml | 1 + docs/contributing/model/adding_tts_model.md | 584 ++++++++++++++++++++ 2 files changed, 585 insertions(+) create mode 100644 docs/contributing/model/adding_tts_model.md diff --git a/docs/.nav.yml b/docs/.nav.yml index d90c7e29a..4aa8fee43 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -66,6 +66,7 @@ nav: - Model Implementation: - contributing/model/README.md - contributing/model/adding_omni_model.md + - contributing/model/adding_tts_model.md - contributing/model/adding_diffusion_model.md - CI: contributing/ci - Design Documents: diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md new file mode 100644 index 000000000..6f642c7cd --- /dev/null +++ b/docs/contributing/model/adding_tts_model.md @@ -0,0 +1,584 @@ +# 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. + +## 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) + +## Overview + +vLLM-Omni supports TTS models as multi-stage pipelines where each stage runs independently +and can be placed on different devices. Qwen3-TTS has two stages: + +| Stage | Name | Input | Output | +|-------|------|-------|--------| +| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes | +| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform | + +Each stage is a separate model class configured independently via YAML. The two stages +are connected by the `async_chunk` framework, which enables inter-stage streaming for +low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk_design.md)). + +### Without async_chunk (batch mode) + +Stage 0 runs to completion before Stage 1 starts, resulting in long first-packet latency: + +```mermaid +flowchart TB + subgraph stage0["Stage 0: AR Code Predictor"] + direction LR + P[Prefill] --> D1[Decode 1] + D1 --> D2[Decode 2] + D2 --> Dots1["..."] + Dots1 --> DN[Decode N] + end + + subgraph stage1["Stage 1: Code2Wav"] + direction LR + DEC[Decode all codes at once] + end + + stage0 -- "all N codes" --> stage1 + stage1 --> FPL["First Packet Latency = Stage 0 + Stage 1"] + + style stage0 fill:#dae8fc,stroke:#6c8ebf + style stage1 fill:#f8d7c8,stroke:#d4856a + style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px +``` + +### With async_chunk (streaming mode) + +Stage 0 sends codec codes to Stage 1 every `chunk_size=25` tokens. Stage 1 begins decoding +immediately, reducing first-packet latency from the full AR time to just the first chunk: + +```mermaid +flowchart TB + subgraph stage0["Stage 0: Code Predictor (AR)"] + direction LR + P[Prefill] --> D1["Decode 1-25"] + D1 --> D2["Decode 26-50"] + D2 --> Dots1["..."] + Dots1 --> DN["Decode N"] + end + + subgraph stage1["Stage 1: Code2Wav"] + direction LR + C1["Chunk 1\n(25 frames)"] --> C2["Chunk 2\n(context + 25)"] + C2 --> Dots2["..."] + Dots2 --> CN["Final chunk"] + end + + D1 -. "chunk 1 (25 codes)" .-> C1 + D2 -. "chunk 2 (context + 25)" .-> C2 + DN -. "final" .-> CN + + stage0 --> FPL["⏱ First Packet Latency = Prefill + 25 decode steps only"] + + style stage0 fill:#dae8fc,stroke:#6c8ebf + style stage1 fill:#e8d4f8,stroke:#8a6cad + style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px +``` + +Key parameters: `chunk_size=25`, `left_context_size=25` (validated defaults from Qwen3-TTS +and Qwen3-Omni). + +## Directory Structure + +When adding a new TTS model, create the following structure: + +``` +vllm_omni/model_executor/models/ + your_model_name/ + __init__.py + your_model.py # Unified class (stage dispatch) + your_model_ar_stage.py # Stage 0: AR stage + your_model_decoder.py # Stage 1: audio decoder + +vllm_omni/model_executor/stage_input_processors/ + your_model_name.py # Stage 0 -> Stage 1 transition + +vllm_omni/model_executor/stage_configs/ + your_model_name.yaml # Batch mode config + your_model_name_async_chunk.yaml # Streaming mode config +``` + +**Qwen3-TTS reference files:** + +| File | Purpose | +|------|---------| +| `models/qwen3_tts/qwen3_tts.py` | Unified model class | +| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR | +| `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder | +| `stage_configs/qwen3_tts.yaml` | Stage config (async_chunk enabled) | +| `stage_configs/qwen3_tts_batch.yaml` | Batch mode config | +| `stage_input_processors/qwen3_tts.py` | Stage transition processors | + +## Step-by-Step Implementation + +### Step 1: Implement Stage 0 - AR Stage + +Stage 0 is the autoregressive stage that generates intermediate audio representations. +**It must use vLLM's native decoder layers with fused ops and PagedAttention** for the LLM +backbone - this is the primary source of speedup over HuggingFace inference. + +#### 1.1 Use vLLM Decoder Layers Directly + +Build your transformer layers from the corresponding vLLM decoder layer class (e.g. +`Qwen3DecoderLayer` for Qwen3-based backbones, or the equivalent for LLaMA, Qwen2, etc.). +Do not wrap the HuggingFace model directly - that bypasses PagedAttention and fused kernels. + +```python +# your_model_ar_stage.py + +from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer + +class YourTTSARStage(nn.Module): + + def __init__(self, config, vllm_config, prefix): + self.layers = nn.ModuleList([ + Qwen3DecoderLayer( + config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}" + ) + for i in range(config.num_hidden_layers) + ]) + self.lm_head = ParallelLMHead(config.codec_size, config.hidden_size) +``` + +See `qwen3_tts_code_predictor_vllm.py` for the full implementation. + +#### 1.2 Forward Pass + +Implement `forward()` to return an `OmniOutput` with intermediate data for Stage 1: + +```python +def forward(self, input_ids, positions, intermediate_tensors=None, + inputs_embeds=None, **kwargs) -> OmniOutput: + hidden_states = self.run_layers(input_ids, positions, intermediate_tensors, inputs_embeds) + logits = self.lm_head(hidden_states) + + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs={ + "audio_codes": self.extract_codes(logits), + }, + ) +``` + +The keys in `multimodal_outputs` are what your stage input processor will read to build +Stage 1 inputs. + +#### 1.3 Weight Loading with Fused QKV + +When using vLLM's fused `QKVParallelLinear`, pack the HF `q_proj`/`k_proj`/`v_proj` weights +into `qkv_proj` using `stacked_params_mapping`. See the `load_weights()` method in +`qwen3_tts_code_predictor_vllm.py` for the standard pattern - it can be reused as-is +for any Qwen-family backbone. + +#### 1.4 Custom Stop Condition (if needed) + +Some TTS models use a learned stop head rather than an EOS token. If your model does this, +implement it inside `sample()`: + +```python +def sample(self, logits, sampling_metadata) -> SamplerOutput | None: + output = self.sampler(logits, sampling_metadata) + if self._stop_head_fired(): + output = mark_as_finished(output) + return output +``` + +### Step 2: Implement Stage 1 - Decoder + +Stage 1 decodes Stage 0 output into audio. It runs outside the scheduler (no PagedAttention +needed). Implement `chunked_decode_streaming()` to support async_chunk streaming: + +```python +# your_model_decoder.py + +class YourTTSDecoder(nn.Module): + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + # Initialize your audio decoder (SpeechTokenizer, HiFiGAN, etc.) + + def forward(self, codes: torch.Tensor, **kwargs) -> torch.Tensor: + return self.decoder(codes) + + def chunked_decode_streaming(self, codes, chunk_size=25, + left_context_size=25) -> torch.Tensor: + """Decode with a sliding context window for smooth chunk boundaries.""" + end_index = codes.shape[-1] + context_size = 0 if end_index <= chunk_size else left_context_size + wav_chunk = self(codes) + # Trim left context to avoid duplicate audio + return wav_chunk[..., context_size * self.total_upsample:] +``` + +### Step 3: Implement the Unified Model Class + +The unified class dispatches to the correct stage based on `model_stage` in the config: + +```python +# your_model.py + +class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP): + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.model_stage = vllm_config.model_config.model_stage + + if self.model_stage == "ar_stage": + ar_vllm_config = vllm_config.with_hf_config( + vllm_config.model_config.hf_config.ar_config, + architectures=["YourTTSARStageForConditionalGeneration"], + ) + self.ar_stage = init_vllm_registered_model( + vllm_config=ar_vllm_config, + prefix=maybe_prefix(prefix, "ar"), + hf_config=ar_vllm_config.model_config.hf_config, + architectures=["YourTTSARStageForConditionalGeneration"], + ) + self.model = self.ar_stage + + elif self.model_stage == "decoder": + self.decoder = YourTTSDecoder(vllm_config=vllm_config, prefix=prefix) + self.model = self.decoder +``` + +### Step 4: Create `__init__.py` + +```python +# vllm_omni/model_executor/models/your_model_name/__init__.py +from .your_model import YourTTSModelForConditionalGeneration + +__all__ = ["YourTTSModelForConditionalGeneration"] +``` + +## Key Components + +### Model Interfaces + +Your unified model class should implement the appropriate interfaces: + +- **`SupportsPP`**: Required for pipeline parallelism support (all models should implement this) +- **`SupportsMultiModal`**: Only if your model accepts multimodal inputs (e.g. reference audio for voice cloning) + +### Output Format + +Use `OmniOutput` so the orchestrator can route intermediate data between stages: + +```python +from vllm_omni.model_executor.models.output_templates import OmniOutput + +return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs={ + "audio_codes": codec_codes, + }, +) +``` + +### Weight Loading from a Single Checkpoint + +If both stages load from one checkpoint, separate them by prefix in the unified class: + +```python +def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + ar_weights, decoder_weights = [], [] + for name, tensor in weights: + if name.startswith("decoder."): + decoder_weights.append((name, tensor)) + else: + ar_weights.append((name, tensor)) + + if self.model_stage == "ar_stage": + return self.ar_stage.load_weights(ar_weights) + elif self.model_stage == "decoder": + return self.decoder.load_weights(decoder_weights) +``` + +## Model Registration + +Register all stage classes in `vllm_omni/model_executor/models/registry.py`: + +```python +_OMNI_MODELS = { + # (package_name, module_name, class_name) + "YourTTSModelForConditionalGeneration": ( + "your_model_name", "your_model", + "YourTTSModelForConditionalGeneration", + ), + "YourTTSARStageForConditionalGeneration": ( + "your_model_name", "your_model_ar_stage", + "YourTTSARStageForConditionalGeneration", + ), + "YourTTSDecoder": ( + "your_model_name", "your_model_decoder", + "YourTTSDecoder", + ), +} +``` + +The registry uses lazy loading - model classes are only imported when needed. + +## Stage Configuration + +Each stage has a `worker_type` that determines how it is scheduled: + +- `worker_type: ar` - autoregressive stage, uses `OmniARScheduler` with PagedAttention +- `worker_type: generation` - non-AR stage (e.g. decoder), uses `OmniGenerationScheduler` + +Key configuration fields: + +| Field | Description | +|-------|-------------| +| `model_stage` | Which stage to initialize (`ar_stage`, `decoder`, etc.) | +| `model_arch` | Architecture name, must match `registry.py` | +| `engine_input_source` | List of upstream stage IDs that provide input (e.g. `[0]`) | +| `engine_output_type` | Output type: `latent` for intermediate, `audio` for final | +| `custom_process_next_stage_input_func` | Async chunk processor function path (streaming only) | +| `final_output` | Whether this stage produces the final user-facing output | +| `final_output_type` | Type of final output (`audio`, `text`, etc.) | + +### Batch mode + +```yaml +# stage_configs/your_model_name.yaml + +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 64 + engine_args: + model_stage: ar_stage + model_arch: YourTTSModelForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + engine_output_type: latent + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 2048 + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + engine_args: + model_stage: decoder + model_arch: YourTTSModelForConditionalGeneration + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + engine_output_type: audio + engine_input_source: [0] + final_output: true + final_output_type: audio +``` + +### Streaming mode (async_chunk) + +Add `async_chunk: true` at the top level and specify `custom_process_next_stage_input_func` +on Stage 0 to define how intermediate outputs are chunked and forwarded: + +```yaml +# stage_configs/your_model_name_async_chunk.yaml + +async_chunk: true + +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 64 + engine_args: + model_stage: ar_stage + model_arch: YourTTSModelForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + engine_output_type: latent + custom_process_next_stage_input_func: > + vllm_omni.model_executor.stage_input_processors.your_model_name.ar2decoder_async_chunk + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 2048 + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + engine_args: + model_stage: decoder + model_arch: YourTTSModelForConditionalGeneration + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + engine_output_type: audio + engine_input_source: [0] + final_output: true + final_output_type: audio +``` + +## Stage Input Processors + +Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours in +`vllm_omni/model_executor/stage_input_processors/your_model_name.py`. + +See `stage_input_processors/qwen3_tts.py` for the full reference implementation. + +### Data structures + +Understanding what's available in stage outputs: + +- `stage_list[source_id].engine_outputs` - list of `EngineCoreOutput` objects +- Each `EngineCoreOutput` has `outputs` - list of `RequestOutput` objects +- Each `RequestOutput` has: + - `token_ids` - generated token IDs + - `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs` + - `prompt_token_ids` - original prompt token IDs + +### Batch mode (non-streaming) + +Collects all Stage 0 outputs and forwards them to Stage 1 in one shot: + +```python +def ar2decoder( + stage_list: list[Any], + engine_input_source: list[int], + prompt: OmniTokensPrompt | TextPrompt | None = None, + requires_multimodal_data: bool = False, +) -> list[OmniTokensPrompt]: + source_id = engine_input_source[0] + decoder_inputs = [] + + for output in stage_list[source_id].engine_outputs: + result = output.outputs[0] + codes = result.multimodal_output["audio_codes"].cpu() + decoder_inputs.append( + OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist()) + ) + + return decoder_inputs +``` + +### Streaming mode (async_chunk) + +Buffers Stage 0 outputs and forwards a chunk to Stage 1 once `chunk_size` frames +have accumulated. The function signature follows the `OmniChunkTransferAdapter` protocol: + +```python +def ar2decoder_async_chunk( + transfer_manager: Any, + pooling_output: dict[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> dict[str, Any] | None: + """Forward chunks of AR output to the decoder stage.""" + request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + + # Extract and buffer the latest frame + if isinstance(pooling_output, dict): + frame = extract_frame(pooling_output) + if frame is not None: + transfer_manager.code_prompt_token_ids[request_id].append( + frame.cpu().tolist() + ) + elif not finished: + return None + + # Read chunk config from connector + chunk_size = 25 + left_context_size = 25 + + length = len(transfer_manager.code_prompt_token_ids[request_id]) + if length <= 0: + if finished: + return {"codes": [], "finished": torch.tensor(True, dtype=torch.bool)} + return None + + # Wait until a full chunk is ready (or request is finished) + chunk_length = length % chunk_size + if chunk_length != 0 and not finished: + return None + + # Build context window: left_context + chunk + context_length = chunk_length if chunk_length != 0 else chunk_size + end_index = min(length, left_context_size + context_length) + window = transfer_manager.code_prompt_token_ids[request_id][-end_index:] + + return { + "codes": torch.tensor(window).transpose(0, 1).reshape(-1).tolist(), + "left_context_size": max(0, int(end_index - context_length)), + "finished": torch.tensor(finished, dtype=torch.bool), + } +``` + +Key points: +- `transfer_manager` is the `OmniChunkTransferAdapter` that owns the chunk lifecycle +- Each call appends one AR decode step's output; a chunk is emitted every `chunk_size` steps +- The final (possibly partial) chunk is flushed when `is_finished` is true +- `left_context_size` frames of overlap are included for smooth audio boundaries + +## Testing + +For general testing conventions, see [tests_style.md](../ci/tests_style.md). + +Recommended test cases for a new TTS model: + +1. **Single request** - verify waveform output shape and sample rate +2. **Batched requests** - verify each request in the batch finishes independently +3. **async_chunk streaming** - verify audio chunks arrive incrementally and decode correctly +4. **Speaker conditioning** (if applicable) - verify different speaker inputs produce different outputs + +Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py` + +## Adding a Model Recipe + +After implementing and testing your model, add a model recipe to the +[vllm-project/recipes](https://github.com/vllm-project/recipes) repository so users can +get started quickly. See [Adding an Omni-Modality Model](./adding_omni_model.md#adding-a-model-recipe) +for the expected format. + +## Summary + +Adding a TTS model to vLLM-Omni involves: + +1. **Create model directory** with AR stage, decoder stage, and unified class +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) + +### Qwen3-TTS Reference Files + +| File | Purpose | +|------|---------| +| `models/qwen3_tts/qwen3_tts.py` | Unified model class | +| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops | +| `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` | +| `stage_configs/qwen3_tts.yaml` | Stage configuration | +| `stage_input_processors/qwen3_tts.py` | Stage transition processors | + +For more information, see: + +- [Architecture Overview](../../design/architecture_overview.md) +- [Async Chunk Design](../../design/feature/async_chunk_design.md) +- [Stage Configuration Guide](../../configuration/stage_configs.md) From b7fcc9d07e07360a8a29bc2e6c981a256289161c Mon Sep 17 00:00:00 2001 From: erfgss <97771661+erfgss@users.noreply.github.com> Date: Fri, 6 Mar 2026 15:39:19 +0800 Subject: [PATCH 8/9] [Profile] Adding metrics for Diffusion/DiT Single diffusion Pipeline (#668) Signed-off-by: Chen Yang <2082464740@qq.com> Signed-off-by: erfgss <97771661+erfgss@users.noreply.github.com> --- docs/contributing/metrics.md | 35 +++++++++- .../image_to_image/image_edit.py | 2 + .../image_to_image/image_edit.py | 6 ++ .../image_to_video/image_to_video.py | 6 ++ .../text_to_image/text_to_image.py | 6 ++ .../text_to_video/text_to_video.py | 6 ++ vllm_omni/diffusion/diffusion_engine.py | 15 ++++- vllm_omni/entrypoints/omni.py | 9 +++ vllm_omni/metrics/stats.py | 64 +++++++++++++++---- vllm_omni/outputs.py | 1 - 10 files changed, 132 insertions(+), 18 deletions(-) diff --git a/docs/contributing/metrics.md b/docs/contributing/metrics.md index 92dcd92cc..bc0f13efd 100644 --- a/docs/contributing/metrics.md +++ b/docs/contributing/metrics.md @@ -69,8 +69,20 @@ With `--log-stats` enabled, the server will output detailed metrics logs after e | rx_decode_time_ms | 111.865 | 31.706 | | in_flight_time_ms | 2.015 | 2.819 | +### DiffusionStats -These logs include: +| Field | Value | +|---------------------------------|------------| +| batch_id | 1 | +| batch_size | 1 | +| diffusion_engine_exec_time_ms | 16,363.419 | +| diffusion_engine_total_time_ms | 16,234.849 | +| image_num | 1.000 | +| num_inference_steps | 50.000 | +| postprocess_time_ms | 67.685 | +| preprocess_time_ms | 60.106 | +| resolution | 640.000 | +| stage_gen_time_ms | 16,363.714 | - **Overall summary**: total requests, wall time, average tokens/sec, etc. @@ -80,6 +92,9 @@ These logs include: - **Transfer table**: data transfer and timing for each edge. +- **DiffusionStats**:preprocessing time, model execution time, postprocessing time, generated images, resolution, and inference steps. + + You can use these logs to monitor system health, debug performance, and analyze request-level metrics as described above. @@ -130,6 +145,11 @@ For **online inference** (serving mode), the summary is always per-request. `e2e | `image_num` | Number of images generated (for diffusion/image stages). | | `resolution` | Image resolution (for diffusion/image stages). | | `postprocess_time_ms` | Diffusion/image: post-processing time in ms. | +| `preprocess_time_ms` | Diffusion/image: pre-processing time in ms (before model execution). | +| `diffusion_engine_exec_time_ms` | Diffusion/image: model execution time in ms, excluding pre-processing and post-processing. | +| `diffusion_engine_total_time_ms` | Diffusion/image: total diffusion engine time in ms, including pre-processing, model execution, and post-processing. | +| `num_inference_steps` | Number of inference steps for diffusion generation. | + --- @@ -143,6 +163,8 @@ For **online inference** (serving mode), the summary is always per-request. `e2e | `in_flight_time_ms` | In-flight time in ms. | + + ### Expectation of the Numbers (Verification) **Formulas:** @@ -171,3 +193,14 @@ For each edge: - 1->2: tx_time_ms (**18.790**) + rx_decode_time_ms (**31.706**) + in_flight_time_ms (**2.819**) = **53.315** 192.581 + 53.315 = **245.896** = transfers_total_time_ms, which matches the calculation (difference is due to rounding) + +**diffusion_engine_total_time_ms** + +- preprocess_time_ms: Time spent on input preprocessing +- diffusion_engine_exec_time_ms: Pure model execution time +- postprocess_time_ms: Time spent on output postprocessing +- diffusion_engine_total_time_ms: Total time spent in the diffusion engine + +so `diffusion_engine_total_time_ms` = `diffusion_engine_exec_time_ms` + `preprocess_time_ms` + `postprocess_time_ms` + = **16,363.419** + **60.106** + **67.685** + = **16,491.210** diff --git a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py index 8ad5cbf9b..6af06fa19 100644 --- a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py +++ b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py @@ -99,6 +99,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--vae_use_slicing", action="store_true") parser.add_argument("--vae_use_tiling", action="store_true") parser.add_argument("--enable-cpu-offload", action="store_true") + parser.add_argument("--log-stats", "--log_stats", dest="log_stats", action="store_true", default=False) return parser.parse_args() @@ -155,6 +156,7 @@ async def main(): cache_config=cache_config, parallel_config=parallel_config, enforce_eager=args.enforce_eager, + log_stats=args.log_stats, enable_cpu_offload=args.enable_cpu_offload, diffusion_load_format="dummy", custom_pipeline_args={"pipeline_class": "custom_pipeline.CustomPipeline"}, diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 81fe24e21..305a6d54a 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -320,6 +320,11 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) + parser.add_argument( + "--log-stats", + action="store_true", + help="Enable vLLM-Omni statistics logging.", + ) return parser.parse_args() @@ -383,6 +388,7 @@ def main(): parallel_config=parallel_config, enforce_eager=args.enforce_eager, enable_cpu_offload=args.enable_cpu_offload, + log_stats=args.log_stats, ) print("Pipeline loaded") diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index cdcabf48c..58da640d9 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -176,6 +176,11 @@ def parse_args() -> argparse.Namespace: "Default 1 means pure sharding (no replication). " ), ) + parser.add_argument( + "--log-stats", + action="store_true", + help="Enable vLLM-Omni statistics logging.", + ) return parser.parse_args() @@ -278,6 +283,7 @@ def main(): enable_cpu_offload=args.enable_cpu_offload, parallel_config=parallel_config, enforce_eager=args.enforce_eager, + log_stats=args.log_stats, model_class_name=model_class_name, cache_backend=args.cache_backend, cache_config=cache_config, diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 382971606..15cbbbf72 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -94,6 +94,11 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable cache-dit summary logging after diffusion forward passes.", ) + parser.add_argument( + "--log-stats", + action="store_true", + help="Enable vLLM-Omni statistics logging.", + ) parser.add_argument( "--ulysses-degree", type=int, @@ -299,6 +304,7 @@ def main(): "parallel_config": parallel_config, "enforce_eager": args.enforce_eager, "enable_cpu_offload": args.enable_cpu_offload, + "log_stats": args.log_stats, **lora_args, **quant_kwargs, } diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index c1efd342e..857530045 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -127,6 +127,11 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for VAE patch/tile parallelism (decode).", ) + parser.add_argument( + "--log-stats", + action="store_true", + help="Enable vLLM-Omni statistics logging.", + ) return parser.parse_args() @@ -176,6 +181,7 @@ def main(): enable_cpu_offload=args.enable_cpu_offload, parallel_config=parallel_config, enforce_eager=args.enforce_eager, + log_stats=args.log_stats, cache_backend=args.cache_backend, cache_config=cache_config, ) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 066ba2bb8..e5ace3b28 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -67,14 +67,19 @@ class DiffusionEngine: raise e def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]: + diffusion_engine_start_time = time.time() # Apply pre-processing if available + preprocess_time = 0.0 if self.pre_process_func is not None: preprocess_start_time = time.time() request = self.pre_process_func(request) preprocess_time = time.time() - preprocess_start_time logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds") + exec_start_time = time.time() output = self.add_req_and_wait_for_response(request) + exec_total_time = time.time() - exec_start_time + if output.error: raise Exception(f"{output.error}") logger.info("Generation completed successfully.") @@ -107,14 +112,20 @@ class DiffusionEngine: outputs = [outputs] if outputs is not None else [] metrics = { + "preprocess_time_ms": preprocess_time * 1000, + "diffusion_engine_exec_time_ms": (time.time() - diffusion_engine_start_time) * 1000, + "diffusion_engine_total_time_ms": exec_total_time * 1000, "image_num": int(request.sampling_params.num_outputs_per_prompt), "resolution": int(request.sampling_params.resolution), - "postprocess_time_ms": postprocess_time * 1000, } + if self.pre_process_func is not None: metrics["preprocessing_time_ms"] = preprocess_time * 1000 # Handle single request or multiple requests + metrics["postprocess_time_ms"] = postprocess_time * 1000 + metrics["num_inference_steps"] = int(request.sampling_params.num_inference_steps) + if len(request.prompts) == 1: # Single request: return single OmniRequestOutput prompt = request.prompts[0] @@ -201,7 +212,7 @@ class DiffusionEngine: ), ) - return results + return results @staticmethod def make_engine(config: OmniDiffusionConfig) -> "DiffusionEngine": diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 9ecc5f762..dddcb12b5 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -1135,6 +1135,15 @@ class Omni(OmniBase): final_output_type=stage.final_output_type, # type: ignore[attr-defined] request_output=engine_outputs, ) + try: + if stage.final_output_type == "text" or metrics.log_stats: + output_to_yield.metrics = metrics.build_output_metrics(stage_id, req_id) + except Exception as e: + # Make metrics contract explicit on failure. + output_to_yield.metrics = {} + logger.exception( + f"[{self._name}] Failed to attach output metrics for req {req_id} at stage {stage_id}: {e}", + ) # Record audio generated frames (only when finished) try: diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 46fec6fcc..068b4a30b 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -39,7 +39,7 @@ class StageRequestStats: final_output_type: str | None = None request_id: str | None = None postprocess_time_ms: float = 0.0 - diffusion_metrics: dict[str, int] = None + diffusion_metrics: dict[str, float] = None audio_generated_frames: int = 0 @property @@ -293,23 +293,15 @@ class OrchestratorAggregator: return # 4. Finished with output: assign text metrics if available - output_to_yield.metrics = {} - stage_event = next( - (evt for evt in reversed(self.stage_events.get(req_id, [])) if evt.stage_id == stage_id), - None, - ) - if stage_event is not None and stage_event.final_output_type == "text": - output_to_yield.metrics = { - "num_tokens_in": stage_event.num_tokens_in, - "num_tokens_out": stage_event.num_tokens_out, - "stage_id": stage_event.stage_id, - "final_output_type": stage_event.final_output_type, - } + output_to_yield.metrics = self.build_output_metrics(stage_id, req_id) # 5. Finished: record audio generated frames self.record_audio_generated_frames(output_to_yield, stage_id, req_id) except Exception: + if output_to_yield is not None: + # Make metrics contract explicit on failure. + output_to_yield.metrics = {} logger.exception( "Failed to process metrics for stage %s, req %s", stage_id, @@ -329,12 +321,56 @@ class OrchestratorAggregator: stats.request_id = req_id stats.final_output_type = final_output_type stats.diffusion_metrics = ( - {k: int(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()} + {k: float(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()} if req_id in self.diffusion_metrics else None ) return stats + def _get_stage_event(self, stage_id: int, req_id: Any) -> StageRequestStats | None: + rid_key = str(req_id) + for evt in reversed(self.stage_events.get(rid_key, [])): + if evt.stage_id == stage_id: + return evt + return None + + def _collect_diffusion_metrics(self, req_id: Any) -> dict[str, float]: + """Aggregate diffusion metrics across all stages for a request.""" + rid_key = str(req_id) + merged: dict[str, float] = {} + for evt in self.stage_events.get(rid_key, []): + if not evt.diffusion_metrics: + continue + for key, value in evt.diffusion_metrics.items(): + merged[key] = merged.get(key, 0.0) + float(value) + return merged + + def build_output_metrics(self, stage_id: int, req_id: Any) -> dict[str, Any]: + stage_event = self._get_stage_event(stage_id, req_id) + if stage_event is None: + return {} + + merged: dict[str, Any] = {} + + if stage_event.final_output_type == "text": + merged.update( + { + "num_tokens_in": stage_event.num_tokens_in, + "num_tokens_out": stage_event.num_tokens_out, + } + ) + + if self.log_stats: + diffusion_metrics = self._collect_diffusion_metrics(req_id) + if diffusion_metrics: + merged.update(diffusion_metrics) + + if merged: + merged["stage_id"] = stage_event.stage_id + merged["final_output_type"] = stage_event.final_output_type + + return merged + def on_stage_metrics( self, stage_id: int, diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 5bbd27a1d..af14dd0e6 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -248,7 +248,6 @@ class OmniRequestOutput: f"images={images_repr}", f"prompt={self.prompt!r}", f"latents={self.latents}", - f"metrics={self.metrics}", f"multimodal_output={self._multimodal_output}", ] From 30191ed112d457fdabbe9ab648c1bb45838b525a Mon Sep 17 00:00:00 2001 From: Didan Deng <33117903+wtomin@users.noreply.github.com> Date: Sat, 7 Mar 2026 01:45:02 +0800 Subject: [PATCH 9/9] [Feature]: Add vae-patch-parallel CLI argument in online serving (#1716) Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> --- .../examples/online_serving/text_to_image.md | 23 +++++++++++++++++++ vllm_omni/entrypoints/async_omni.py | 2 ++ vllm_omni/entrypoints/cli/serve.py | 9 ++++++++ 3 files changed, 34 insertions(+) diff --git a/docs/user_guide/examples/online_serving/text_to_image.md b/docs/user_guide/examples/online_serving/text_to_image.md index be8c7d894..793129488 100644 --- a/docs/user_guide/examples/online_serving/text_to_image.md +++ b/docs/user_guide/examples/online_serving/text_to_image.md @@ -23,6 +23,29 @@ Or use the startup script: bash run_server.sh ``` +### Start with Parallelism Acceleration + +Enable Tensor Parallelism and VAE Patch Parallelism for faster inference: + +```bash +# With Tensor Parallelism (requires >= 2 GPUs) +vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2 + +# With Tensor Parallelism and VAE Patch Parallelism (requires >= 2 GPUs) +vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2 --vae-patch-parallel-size 2 --vae-use-tiling + +# With Sequence Parallelism (Ulysses-SP, requires >= 2 GPUs) +vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 + +# With Ring-Attention (requires >= 2 GPUs) +vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2 + +# Combined: Ulysses + Ring (requires >= 4 GPUs) +vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2 +``` + +For more details on parallelism acceleration, see the [Parallelism Acceleration Guide](../../diffusion/parallelism_acceleration.md). + ## API Calls ### Method 1: Using curl diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 6939987b0..e393f31be 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -141,6 +141,7 @@ class AsyncOmni(OmniBase): sequence_parallel_size = kwargs.get("sequence_parallel_size") tensor_parallel_size = kwargs.get("tensor_parallel_size") or 1 cfg_parallel_size = kwargs.get("cfg_parallel_size") or 1 + vae_patch_parallel_size = kwargs.get("vae_patch_parallel_size") or 1 use_hsdp = kwargs.get("use_hsdp", False) hsdp_shard_size = kwargs.get("hsdp_shard_size", -1) hsdp_replicate_size = kwargs.get("hsdp_replicate_size", 1) @@ -165,6 +166,7 @@ class AsyncOmni(OmniBase): ulysses_degree=ulysses_degree, ring_degree=ring_degree, cfg_parallel_size=cfg_parallel_size, + vae_patch_parallel_size=vae_patch_parallel_size, use_hsdp=use_hsdp, hsdp_shard_size=hsdp_shard_size, hsdp_replicate_size=hsdp_replicate_size, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index a580f9536..7365e73bb 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -219,6 +219,7 @@ class OmniServeCommand(CLISubcommand): ) omni_config_group.add_argument( "--ring", + "--ring-degree", dest="ring_degree", type=int, default=None, @@ -335,6 +336,14 @@ class OmniServeCommand(CLISubcommand): help="Number of devices for CFG parallel computation for diffusion models. " "Equivalent to setting DiffusionParallelConfig.cfg_parallel_size.", ) + omni_config_group.add_argument( + "--vae-patch-parallel-size", + type=int, + default=1, + help="VAE Patch Parallelism degree for diffusion models. " + "Distributes VAE decode workload across multiple ranks by splitting the latent spatially. " + "Equivalent to setting DiffusionParallelConfig.vae_patch_parallel_size.", + ) # Default sampling parameters omni_config_group.add_argument(