From 4e1c13236ee26e7bb72da6e567256b87058e7640 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Thu, 4 Jun 2026 09:08:13 +0800 Subject: [PATCH] [Perf][Qwen3-Omni]Optimize TTFP using initial_codec_chunk_frames (#4054) Signed-off-by: amy-why-3459 --- vllm_omni/deploy/qwen3_omni_moe.yaml | 1 + .../deploy/qwen3_omni_moe_multi_replicas.yaml | 1 + .../stage_input_processors/qwen3_omni.py | 64 +++++++++---------- 3 files changed, 34 insertions(+), 32 deletions(-) diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index 13d560145..3bfa6042a 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/vllm_omni/deploy/qwen3_omni_moe.yaml @@ -18,6 +18,7 @@ connectors: connector_of_shared_memory: name: SharedMemoryConnector extra: + initial_codec_chunk_frames: 4 codec_chunk_frames: 25 codec_left_context_frames: 25 diff --git a/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml b/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml index 390698e94..4f94fea60 100644 --- a/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml +++ b/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml @@ -18,6 +18,7 @@ connectors: connector_of_shared_memory: name: SharedMemoryConnector extra: + initial_codec_chunk_frames: 4 codec_chunk_frames: 25 codec_left_context_frames: 25 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index aa98c2141..c2b33368d 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -817,56 +817,56 @@ def talker2code2wav_async_chunk( if code_predictor_codes is None: return None + if code_predictor_codes.numel() == 0: + return None + + if not code_predictor_codes.any(): + return None + connector = getattr(transfer_manager, "connector", None) raw_cfg = getattr(connector, "config", {}) or {} cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} chunk_size_config = int(cfg.get("codec_chunk_frames", 25)) left_context_size_config = int(cfg.get("codec_left_context_frames", 25)) + configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames") or 0) - if code_predictor_codes is None: - return None - if isinstance(code_predictor_codes, torch.Tensor): - if code_predictor_codes.numel() == 0: - return None - elif hasattr(code_predictor_codes, "__len__"): - if len(code_predictor_codes) == 0: - return None - - if isinstance(code_predictor_codes, torch.Tensor): - if not code_predictor_codes.any(): - return None - sampling_params = getattr(request, "sampling_params", None) - stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or []) - stop_token_id = getattr(sampling_params, "stop_token_id", None) - if stop_token_id is not None: - stop_token_ids.add(stop_token_id) - first_codebook = int(code_predictor_codes[0, 0].item()) - if first_codebook in stop_token_ids: - logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook) - return None - else: - code_tensor = torch.tensor(code_predictor_codes, dtype=torch.long) - if not code_tensor.any(): - return None - - codec_codes = code_predictor_codes.to(torch.long).transpose(0, 1).cpu().to(torch.long).reshape(-1).tolist() - if sum(codec_codes) == 0: + sampling_params = getattr(request, "sampling_params", None) + stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or []) + stop_token_id = getattr(sampling_params, "stop_token_id", None) + if stop_token_id is not None: + stop_token_ids.add(stop_token_id) + first_codebook = int(code_predictor_codes[0, 0].item()) + if first_codebook in stop_token_ids: + logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook) return None request_id = request.external_req_id - transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) + chunk_id = transfer_manager.put_req_chunk[request_id] + transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes) length = len(transfer_manager.code_prompt_token_ids[request_id]) + if configured_initial_chunk_size > 0: + if chunk_id == 0: + chunk_size_config = configured_initial_chunk_size + else: + length -= configured_initial_chunk_size + chunk_length = length % chunk_size_config if chunk_length != 0 and not is_finished: return None context_length = chunk_length if chunk_length != 0 else chunk_size_config # ensure left context does not exceed available length - left_context_size = max(0, min(length - context_length, left_context_size_config)) - end_index = min(length, left_context_size + context_length) + if configured_initial_chunk_size > 0 and chunk_id == 1: + left_context_size = configured_initial_chunk_size + end_index = length + configured_initial_chunk_size + else: + left_context_size = max(0, min(length - context_length, left_context_size_config)) + end_index = min(length, left_context_size + context_length) - codes = torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]).transpose(0, 1).reshape(-1) + codes = ( + torch.cat(transfer_manager.code_prompt_token_ids[request_id][-end_index:], dim=0).transpose(0, 1).reshape(-1) + ) return OmniPayloadStruct( codes=CodesStruct(audio=codes),