diff --git a/tests/core/test_prefix_cache.py b/tests/core/test_prefix_cache.py index b5d0e96d3..326eb6b34 100644 --- a/tests/core/test_prefix_cache.py +++ b/tests/core/test_prefix_cache.py @@ -13,7 +13,7 @@ DEFAULT_SHAPE = torch.Size([NUM_BLOCKS, BLOCK_SIZE, HIDDEN_SIZE]) class MockInputBatch: - def __init__(self, num_computed_tokens_cpu): + def __init__(self, num_computed_tokens_cpu, block_table=None): self.req_ids = ["req1", "req2"] self.req_id_to_index = {req_id: i for i, req_id in enumerate(self.req_ids)} self.num_computed_tokens_cpu = num_computed_tokens_cpu @@ -21,11 +21,33 @@ class MockInputBatch: # Block table is only mocked for validation of length; # we don't actually need to add valid values here since # we patch the table when testing. - class _DummyBlockTable: - pass + if block_table is None: - self.block_table = _DummyBlockTable() - self.block_table.block_tables = [None] + class _DummyBlockTable: + pass + + self.block_table = _DummyBlockTable() + self.block_table.block_tables = [None] + else: + + class _TensorWrapper: + def __init__(self, tensor): + self.cpu = tensor + + class _BlockGroup: + def __init__(self, tensor): + self.block_table = _TensorWrapper(tensor) + + class _BlockTable: + def __init__(self, tensor): + self._group = _BlockGroup(tensor) + self.block_tables = [self._group.block_table] + + def __getitem__(self, idx): + assert idx == 0 + return self._group + + self.block_table = _BlockTable(block_table) def get_omni_pcache_with_mm_tensors(feat_dims, seq_len) -> OmniTensorPrefixCache: @@ -171,6 +193,94 @@ def test_update_with_multimodal_outputs(feat_dims): assert torch.all(slot_states == new_output) +def test_deferred_multimodal_cache_write_commits_on_completion(): + cache = get_omni_pcache() + mm_key = "codes.audio" + block_table = torch.tensor( + [ + [2, 3, 4, 5], + [6, 7, 8, 9], + ], + dtype=torch.long, + ) + input_batch = MockInputBatch( + num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long), + block_table=block_table, + ) + + first_chunk = torch.arange(8, dtype=torch.long).reshape(4, 2) + cache.update_omni_tensor_prefix_cache( + hidden_states=None, + multimodal_outputs={mm_key: first_chunk}, + num_tokens_unpadded=4, + slot_mapping=torch.arange(8, 12), + skip_mm_cache_keys={mm_key}, + ) + + rows = cache.mm_outputs_cache[mm_key].view(-1, 2) + assert torch.all(rows[8:12] == 0) + + cache.stage_deferred_mm_outputs( + query_start_loc=torch.tensor([0, 4], dtype=torch.long), + input_batch=input_batch, + multimodal_outputs={mm_key: first_chunk}, + num_scheduled_tokens={"req1": 4, "req2": 0}, + deferred_mm_cache_keys={mm_key}, + ) + + second_chunk = torch.arange(8, 16, dtype=torch.long).reshape(4, 2) + input_batch.num_computed_tokens_cpu = torch.tensor([8, 0], dtype=torch.long) + cache.stage_deferred_mm_outputs( + query_start_loc=torch.tensor([0, 4], dtype=torch.long), + input_batch=input_batch, + multimodal_outputs={mm_key: second_chunk}, + num_scheduled_tokens={"req1": 4, "req2": 0}, + deferred_mm_cache_keys={mm_key}, + ) + + cache.commit_deferred_mm_outputs({"req1"}, input_batch) + + expected = torch.cat([first_chunk, second_chunk], dim=0) + rows = cache.mm_outputs_cache[mm_key].view(-1, 2) + assert torch.equal(rows[8:16], expected) + + +def test_deferred_multimodal_cache_can_be_merged_on_full_block_hit(): + cache = get_omni_pcache() + mm_key = "codes.audio" + block_table = torch.tensor( + [ + [2, 3, 4, 5], + [6, 7, 8, 9], + ], + dtype=torch.long, + ) + input_batch = MockInputBatch( + num_computed_tokens_cpu=torch.tensor([8, 0], dtype=torch.long), + block_table=block_table, + ) + cached_codes = torch.arange(16, dtype=torch.long).reshape(8, 2) + cache.stage_deferred_mm_outputs( + query_start_loc=torch.tensor([0, 8], dtype=torch.long), + input_batch=input_batch, + multimodal_outputs={mm_key: cached_codes}, + num_scheduled_tokens={"req1": 8, "req2": 0}, + deferred_mm_cache_keys={mm_key}, + ) + cache.commit_deferred_mm_outputs({"req1"}, input_batch) + + new_codes = torch.tensor([[99, 100]], dtype=torch.long) + cache.add_prefix_cached_new_req_id("req1") + merged = cache.get_merged_multimodal_states( + query_start_loc=torch.tensor([0, 1], dtype=torch.long), + input_batch=input_batch, + multimodal_outputs={mm_key: new_codes}, + num_scheduled_tokens={"req1": 1, "req2": 0}, + ) + + assert torch.equal(merged[mm_key]["req1"], torch.cat([cached_codes, new_codes], dim=0)) + + ### Tests for Merging def fake_get_cached_block_ids(self, req_idx, *args, **kwargs): """Fake block table lookup. diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py index fd3fecaa9..d4b8e4a32 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py @@ -1,3 +1,4 @@ +from collections import OrderedDict from types import SimpleNamespace import numpy as np @@ -15,6 +16,8 @@ from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( def _make_minimal_talker(): model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16) + model._ref_audio_artifact_cache_max_entries = 256 + model._ref_audio_artifact_cache = OrderedDict() return model diff --git a/vllm_omni/core/prefix_cache.py b/vllm_omni/core/prefix_cache.py index 69e7346c4..f80074173 100644 --- a/vllm_omni/core/prefix_cache.py +++ b/vllm_omni/core/prefix_cache.py @@ -47,6 +47,7 @@ class OmniTensorPrefixCache: self.mm_outputs_cache = {} self.mm_cache_keys = set() self._new_req_cache_hit_ids: set[str] = set() + self._deferred_mm_outputs: dict[str, dict[str, list[tuple[int, torch.Tensor]]]] = {} def maybe_init_missing_mm_cache_keys(self, multimodal_outputs: dict, seq_len: int): """Given multimodal outputs from executing the model, dynamically @@ -62,7 +63,17 @@ class OmniTensorPrefixCache: determined by the warmup. """ for key, val in multimodal_outputs.items(): - if isinstance(val, torch.Tensor) and val.shape[0] == seq_len and key not in self.mm_cache_keys: + # Only cache per-token feature tensors: 2D+ with first dim == seq_len. + # A 1D tensor of shape (seq_len,) is a broadcast scalar (per-request + # metadata such as ref_code_len / codec_streaming), not per-token data; + # caching it by slot causes a shape mismatch when a later request has a + # different scheduled seq length. + if ( + isinstance(val, torch.Tensor) + and val.ndim >= 2 + and val.shape[0] == seq_len + and key not in self.mm_cache_keys + ): feat_dim = val.shape[-1] self.mm_outputs_cache[key] = self._get_cache_tensor( dtype=val.dtype, @@ -89,6 +100,10 @@ class OmniTensorPrefixCache: """Clears the cache hit IDs to prepare for a new engine step.""" self._new_req_cache_hit_ids.clear() + def has_prefix_cached_new_req_ids(self) -> bool: + """Return True when this step contains a newly scheduled prefix hit.""" + return bool(self._new_req_cache_hit_ids) + @staticmethod def _coerce_to_cpu_tensor(maybe_gpu_tensor: torch.Tensor) -> torch.Tensor: """Convert GPU tensors -> contiguous CPU tensors if needed.""" @@ -101,6 +116,7 @@ class OmniTensorPrefixCache: num_tokens_unpadded: int, slot_mapping: torch.Tensor, num_tokens_padded: int | None = None, + skip_mm_cache_keys: set[str] | None = None, ): """Updates the hidden cache state for the provided hidden states and multimodal outputs. @@ -110,10 +126,12 @@ class OmniTensorPrefixCache: num_tokens_unpadded: Number of tokens without padding slot_mapping: Slot mapping for the input sequence num_tokens_padded: Total number of tokens including padding + skip_mm_cache_keys: Multimodal keys whose CPU cache writes are deferred """ unpadded_slot_mapping = slot_mapping[:num_tokens_unpadded] if num_tokens_padded is None: num_tokens_padded = num_tokens_unpadded + skip_mm_cache_keys = skip_mm_cache_keys or set() if hidden_states is not None: # Slice to unpadded portion before caching @@ -136,6 +154,8 @@ class OmniTensorPrefixCache: for mm_out_key, mm_cache in self.mm_outputs_cache.items(): if mm_out_key in multimodal_outputs: + if mm_out_key in skip_mm_cache_keys: + continue # Slice to unpadded portion before caching mm_state = multimodal_outputs[mm_out_key][:num_tokens_unpadded] mm_state = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state) @@ -143,6 +163,165 @@ class OmniTensorPrefixCache: flat_cache[unpadded_slot_mapping] = mm_state logger.debug("Writing to mm output cache for %s tokens", num_tokens_unpadded) + def stage_deferred_mm_outputs( + self, + query_start_loc: torch.Tensor, + input_batch: InputBatch, + multimodal_outputs: dict[str, torch.Tensor] | None, + num_scheduled_tokens: dict[str, int], + deferred_mm_cache_keys: set[str], + ) -> None: + """Keep GPU multimodal chunks until a request finishes. + + The normal prefix-cache write path copies every cached multimodal tensor + to CPU on every step. For model outputs that are only needed by future + full-block prefix hits, we can keep detached GPU chunks and materialize + the CPU cache once when the request completes. + """ + if not multimodal_outputs or not deferred_mm_cache_keys: + return + + for mm_key in deferred_mm_cache_keys: + mm_state = multimodal_outputs.get(mm_key) + if not isinstance(mm_state, torch.Tensor) or mm_state.ndim < 2: + continue + for req_id in input_batch.req_ids: + req_idx = input_batch.req_id_to_index[req_id] + sched = int(num_scheduled_tokens.get(req_id, 0)) + if sched <= 0: + continue + start = int(query_start_loc[req_idx]) + end = start + sched + if start >= int(mm_state.shape[0]): + continue + end = min(end, int(mm_state.shape[0])) + if end <= start: + continue + + computed_end = int(input_batch.num_computed_tokens_cpu[req_idx]) + token_start = max(0, computed_end - sched) + chunk = mm_state[start:end].detach() + self._deferred_mm_outputs.setdefault(req_id, {}).setdefault(mm_key, []).append((token_start, chunk)) + + def commit_deferred_mm_outputs( + self, + finished_req_ids: set[str] | list[str], + input_batch: InputBatch, + ) -> None: + """Write deferred multimodal chunks into the CPU prefix cache. + + This must run before finished requests are removed from ``input_batch``, + because the block table is needed to map logical token positions to KV + cache slots. + """ + if not finished_req_ids or not self._deferred_mm_outputs: + return + + for req_id in tuple(finished_req_ids): + per_key = self._deferred_mm_outputs.pop(req_id, None) + if not per_key: + continue + req_idx = input_batch.req_id_to_index.get(req_id) + if req_idx is None: + continue + for mm_key, chunks in per_key.items(): + self._commit_deferred_mm_key(req_idx, input_batch, mm_key, chunks) + + def discard_deferred_mm_outputs(self, req_id: str) -> None: + """Drop deferred chunks for requests that leave without a cache commit.""" + self._deferred_mm_outputs.pop(req_id, None) + + def _commit_deferred_mm_key( + self, + req_idx: int, + input_batch: InputBatch, + mm_key: str, + chunks: list[tuple[int, torch.Tensor]], + ) -> None: + if not chunks: + return + + chunks = sorted(chunks, key=lambda item: item[0]) + first_start = chunks[0][0] + expected_start = first_start + tensors: list[torch.Tensor] = [] + for token_start, tensor in chunks: + if tensor.numel() == 0: + continue + if token_start != expected_start: + if tensors: + self._write_deferred_mm_tensor(req_idx, input_batch, mm_key, tensors, first_start) + first_start = token_start + expected_start = token_start + tensors = [] + tensors.append(tensor) + expected_start += int(tensor.shape[0]) + + if tensors: + self._write_deferred_mm_tensor(req_idx, input_batch, mm_key, tensors, first_start) + + def _write_deferred_mm_tensor( + self, + req_idx: int, + input_batch: InputBatch, + mm_key: str, + tensors: list[torch.Tensor], + token_start: int, + ) -> None: + mm_state = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0) + if mm_state.ndim < 2 or mm_state.shape[0] == 0: + return + + if mm_key not in self.mm_cache_keys: + self.mm_outputs_cache[mm_key] = self._get_cache_tensor( + dtype=mm_state.dtype, + hidden_size=int(mm_state.shape[-1]), + ) + self.mm_cache_keys.add(mm_key) + logger.info( + "Initializing deferred multimodal output cache of size %s for key: %s", + list(self.mm_outputs_cache[mm_key].shape), + mm_key, + ) + + slots = self._get_slot_ids_for_token_range( + req_idx=req_idx, + input_batch=input_batch, + token_start=token_start, + num_tokens=int(mm_state.shape[0]), + ) + if slots.numel() == 0: + return + + mm_state_cpu = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state[: slots.numel()]) + mm_cache = self.mm_outputs_cache[mm_key] + flat_cache = mm_cache.view(-1, mm_cache.shape[-1]) + flat_cache[slots] = mm_state_cpu + + def _get_slot_ids_for_token_range( + self, + req_idx: int, + input_batch: InputBatch, + token_start: int, + num_tokens: int, + ) -> torch.Tensor: + if num_tokens <= 0: + return torch.empty((0,), dtype=torch.long) + + block_table = input_batch.block_table[0].block_table.cpu + token_positions = torch.arange(token_start, token_start + num_tokens, dtype=torch.long) + block_offsets = token_positions // self.block_size + max_blocks = int(block_table.shape[1]) + valid = block_offsets < max_blocks + if not bool(valid.all()): + token_positions = token_positions[valid] + block_offsets = block_offsets[valid] + if token_positions.numel() == 0: + return torch.empty((0,), dtype=torch.long) + + block_ids = block_table[req_idx, block_offsets].to(torch.long) + return block_ids * self.block_size + (token_positions % self.block_size) + def _coerce_to_payload_dict( self, element: object, @@ -151,17 +330,25 @@ class OmniTensorPrefixCache: num_scheduled_tokens: dict[str, int], ) -> dict[str, object]: """Build the multimodal passthrough data per request for - the object under consideration. This is identical to the case - for no prefix cache when we tensor does have a first dimension - matching the seq len. + the object under consideration. This mirrors the no-prefix-cache + path: a tensor whose first dimension matches the total scheduled + token count is sliced per request, so 1D (seq_len,) metadata that + is intentionally not cached (e.g. ref_code_len, codec_streaming) + is still split per request instead of leaking the whole batch. """ + total_scheduled_tokens = sum(int(num_scheduled_tokens[r]) for r in input_batch.req_ids) elem_dict = {} for req_id in input_batch.req_ids: req_idx = input_batch.req_id_to_index[req_id] start = query_start_loc[req_idx] end = start + num_scheduled_tokens[req_id] elem_dict[req_id] = to_payload_element( - element, req_idx, start=start, end=end, pass_lists_through=True, seq_len=None + element, + req_idx, + start=start, + end=end, + pass_lists_through=True, + seq_len=total_scheduled_tokens, ) return elem_dict diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 0c976d090..f54cca6ef 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -34,7 +34,12 @@ stages: max_num_seqs: 10 gpu_memory_utilization: 0.3 trust_remote_code: true - enable_prefix_caching: false + # Prefix caching on by default. The HS opt-out (talker only consumes the + # last token's hidden state) and deferred mm-output writes added in this + # PR make per-step bookkeeping negligible, and any text repetition across + # requests (same voice or shared ref_audio) is a net throughput win. + # Disable explicitly if memory headroom is tight on a single-GPU deploy. + enable_prefix_caching: true async_scheduling: true max_num_batched_tokens: 512 max_model_len: 4096 diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 27bb44ad7..ae0c82765 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -107,7 +107,7 @@ def build_engine_core_request_from_tokens( pooling_params=pooling_params, arrival_time=arrival_time, lora_request=getattr(params, "lora_request", None), - cache_salt=None, + cache_salt=prompt.get("cache_salt"), data_parallel_rank=None, prompt_embeds=prompt_embeds, resumable=resumable, diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 7bd27be81..de4ef327c 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -277,6 +277,10 @@ class OmniBase(PDDisaggregationMixin): try: if req_state is None or req_state.metrics is None: return + if self.log_stats: + # Emit per-request orchestrator timing (including e2e_total_ms) + # before dropping request state. + req_state.metrics.build_and_log_summary() except Exception: logger.exception( "[%s] Failed to build/log summary for req=%s", diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 10036978e..7643ed625 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1,5 +1,6 @@ import asyncio import base64 +import hashlib import io import json import math @@ -186,6 +187,55 @@ def _validate_path_within_directory(file_path: Path, directory: Path) -> bool: return False +def _conditioning_cache_salt(request, tts_params: dict | None = None) -> str: + """Stable hash of the real Stage 0 conditioning for the prefix cache. + + The talker's vLLM prompt is placeholder token ids; the real inputs are + rebuilt from text / ref_audio / ref_text into inputs_embeds. vLLM hashes + token ids (folded with cache_salt) for prefix caching, so without a salt + every request collides and a hit could reuse KV from a semantically + different input. Tying the salt to the conditioning keeps a hit safe: + identical conditioning may share the prefix, any difference never does. + + Raw request fields alone are not enough for uploaded voices: the request + only carries the voice *name* (ref_audio/ref_text/task_type are resolved + from stored voice data into ``tts_params``). Delete + re-upload under the + same name leaves every raw field identical, so the resolved conditioning + must also be folded in. ``voice_created_at`` bumps on every (re-)upload, + which uniquely identifies the resolved reference artifact together with + the voice name; the decoded ref_audio array itself need not be hashed. + """ + h = hashlib.sha256() + for part in ( + request.input, + request.task_type, + request.language, + request.voice, + request.ref_text, + request.ref_audio, + request.instructions, + request.x_vector_only_mode, + request.speaker_embedding, + ): + h.update(b"\x00") + if part is not None: + h.update(repr(part).encode("utf-8")) + # Fold resolved conditioning that is auto-derived for uploaded voices and + # absent from the raw request. + for key in ( + "voice_created_at", + "task_type", + "speaker", + "ref_text", + "x_vector_only_mode", + ): + h.update(b"\x00") + value = tts_params.get(key) if tts_params is not None else None + if value is not None: + h.update(repr(value).encode("utf-8")) + return h.hexdigest()[:32] + + class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): _diffusion_mode: bool = False _tts_executor: ThreadPoolExecutor | None = None @@ -1500,6 +1550,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id: str, response_format: str = "pcm", raw_request: Request | None = None, + request_start_s: float | None = None, ): """Generate audio chunks for streaming response. @@ -1520,6 +1571,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): prev_count = 0 sample_rate_val = 24000 first_chunk = True + first_audio_chunk_s: float | None = None + stream_start_s = request_start_s if request_start_s is not None else time.perf_counter() try: async for res in generator: @@ -1571,11 +1624,40 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): stream_format="audio", base64_encode=False, ) + if first_audio_chunk_s is None: + first_audio_chunk_s = time.perf_counter() yield self.create_audio(audio_obj).audio_data + total_ms = (time.perf_counter() - stream_start_s) * 1000.0 + if first_audio_chunk_s is not None: + first_chunk_ms = (first_audio_chunk_s - stream_start_s) * 1000.0 + logger.info( + "[SpeechE2E] request_id=%s stream=true status=ok total_ms=%.2f first_chunk_ms=%.2f", + request_id, + total_ms, + first_chunk_ms, + ) + else: + logger.info( + "[SpeechE2E] request_id=%s stream=true status=ok total_ms=%.2f first_chunk_ms=NA", + request_id, + total_ms, + ) except asyncio.CancelledError: + total_ms = (time.perf_counter() - stream_start_s) * 1000.0 + logger.info( + "[SpeechE2E] request_id=%s stream=true status=cancelled total_ms=%.2f", + request_id, + total_ms, + ) logger.info("Streaming request %s cancelled by client", request_id) raise except EngineDeadError as e: + total_ms = (time.perf_counter() - stream_start_s) * 1000.0 + logger.error( + "[SpeechE2E] request_id=%s stream=true status=engine_dead total_ms=%.2f", + request_id, + total_ms, + ) logger.error( "EngineDeadError during streaming speech for %s: %s", request_id, @@ -1589,6 +1671,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) raise except Exception as e: + total_ms = (time.perf_counter() - stream_start_s) * 1000.0 + logger.exception( + "[SpeechE2E] request_id=%s stream=true status=error total_ms=%.2f error=%s", + request_id, + total_ms, + e, + ) logger.exception("Streaming speech generation failed for %s: %s", request_id, e) raise @@ -2195,6 +2284,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): tts_params["seed"] = [sampling_params_list[0].seed] prompt = tokens_input(prompt_token_ids=[1]) prompt["additional_information"] = tts_params + prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params) else: tts_params = self._build_tts_params(request) # Resolve ref_audio (explicit or auto-set for uploaded voices) @@ -2210,6 +2300,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ph_len = await self._estimate_prompt_len_async(tts_params) prompt = tokens_input(prompt_token_ids=[1] * ph_len) prompt["additional_information"] = tts_params + prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params) else: # Qwen omni models (Qwen3-Omni, Qwen2.5-Omni) use a "talker" # stage whose preprocess requires chat-templated tokens. The @@ -2651,6 +2742,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return error_check_ret request_id = f"speech-{random_uuid()}" + request_start_s = time.perf_counter() if raw_request: raw_request.state.request_metadata = RequestResponseMetadata( request_id=request_id, @@ -2683,20 +2775,58 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id, response_format, raw_request=raw_request, + request_start_s=request_start_s, ), media_type=media_type, ) audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id) + total_ms = (time.perf_counter() - request_start_s) * 1000.0 + logger.info( + "[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d", + request_id, + total_ms, + len(audio_bytes) if isinstance(audio_bytes, (bytes, bytearray)) else len(str(audio_bytes)), + ) return Response(content=audio_bytes, media_type=media_type) except asyncio.CancelledError: + total_ms = (time.perf_counter() - request_start_s) * 1000.0 + logger.info( + "[SpeechE2E] request_id=%s stream=%s status=cancelled total_ms=%.2f", + request_id, + bool(request.stream), + total_ms, + ) return self.create_error_response("Client disconnected") except (EngineGenerateError, EngineDeadError): + total_ms = (time.perf_counter() - request_start_s) * 1000.0 + logger.error( + "[SpeechE2E] request_id=%s stream=%s status=engine_error total_ms=%.2f", + request_id, + bool(request.stream), + total_ms, + ) raise # Propagate to the global Omni exception handler except ValueError as e: + total_ms = (time.perf_counter() - request_start_s) * 1000.0 + logger.warning( + "[SpeechE2E] request_id=%s stream=%s status=bad_request total_ms=%.2f error=%s", + request_id, + bool(request.stream), + total_ms, + e, + ) return self.create_error_response(e) except Exception as e: + total_ms = (time.perf_counter() - request_start_s) * 1000.0 + logger.exception( + "[SpeechE2E] request_id=%s stream=%s status=error total_ms=%.2f error=%s", + request_id, + bool(request.stream), + total_ms, + e, + ) logger.exception("Speech generation failed: %s", e) return self.create_error_response(f"Speech generation failed: {e}") diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index a705a1cd4..470ad34d5 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -2,8 +2,10 @@ from __future__ import annotations import base64 import copy +import hashlib import io import os +from collections import OrderedDict from collections.abc import Callable, Iterable, Mapping from typing import Any from urllib.parse import urlparse @@ -40,6 +42,7 @@ logger = init_logger(__name__) _TRAILING_TEXT_COMPACT_MIN_FRAMES = 64 _PRECOMPUTED_REF_CODE_KEY = "precomputed_ref" _NORMALIZED_REF_AUDIO_KEY = "_qwen3_tts_normalized_ref_audio" +_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key" _PRECOMPUTED_TEXT_IDS_KEY = "_qwen3_tts_text_ids" _PRECOMPUTED_REF_IDS_KEY = "_qwen3_tts_ref_ids" @@ -354,6 +357,17 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): self.have_multimodal_outputs = True self.has_preprocess = True self.has_postprocess = True + # Qwen3-TTS postprocess() only reads hidden_states[-1, :]. On a prefix- + # cache hit, the last hidden state is in the newly computed tail, so + # reconstructing the full cached_prefix + new_tail span is wasted work. + # Opt out of the per-step GPU->CPU hidden-state cache write and merged- + # tensor read; postprocess receives the tail-only slice instead, which + # avoids ~18 ms merge + ~6 ms write per step (Sy0307 profile, #3665). + self.requires_full_prefix_cached_hidden_states = False + # ``codes.audio`` is only needed for future prefix-hit reconstruction + # after a request has produced codec rows. Keep per-step rows on GPU and + # materialize the CPU OmniTensorPrefixCache entry once at completion. + self.deferred_prefix_cache_mm_keys = {"codes.audio"} # Used by OmniGPUModelRunner for the GPU-side MTP fast-path. self.mtp_hidden_size = int(self.talker_config.hidden_size) # OmniGPUModelRunner will store talker_mtp output under this key in @@ -443,6 +457,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): self._speech_tokenizer: Qwen3TTSTokenizer | None = None self._speaker_cache = get_speaker_cache() + self._ref_audio_artifact_cache_max_entries = 256 + self._ref_audio_artifact_cache: OrderedDict[str, dict[str, torch.Tensor | bool]] = OrderedDict() raw_subtalker_sampling = getattr(vllm_config.model_config, "subtalker_sampling_params", None) self._subtalker_sampling_params: dict[str, Any] = ( dict(raw_subtalker_sampling) if isinstance(raw_subtalker_sampling, Mapping) else {} @@ -634,9 +650,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): info_update["meta"]["ref_code_len"] = int(ref_code_len) # Always return a span_len slice; if the scheduled placeholder is longer, pad with tts_pad_embed. # This preserves placeholder/embedding alignment. - offset = 0 - s = 0 - e = span_len + offset = int(info_dict.get("_omni_num_computed_tokens", 0) or 0) + if offset < 0: + offset = 0 + s = max(0, min(offset, int(prompt_embeds_cpu.shape[0]))) + e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0]))) take = prompt_embeds_cpu[s:e] if int(take.shape[0]) < span_len: pad_n = int(span_len - int(take.shape[0])) @@ -1393,9 +1411,50 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # 12Hz: likely [T, Q] or [B, T, Q] if ref_code.ndim == 3: ref_code = ref_code[0] - return ref_code.to(device=next(self.parameters()).device, dtype=torch.long) + out = ref_code.to(device=next(self.parameters()).device, dtype=torch.long) + return out raise ValueError("SpeechTokenizer.encode did not return audio_codes tensor") + @staticmethod + def _make_ref_audio_cache_key(wav: np.ndarray, sr: int) -> str: + wav_f32 = wav.astype(np.float32, copy=False).reshape(-1) + h = hashlib.sha1() + h.update(int(sr).to_bytes(4, byteorder="little", signed=False)) + h.update(int(wav_f32.size).to_bytes(8, byteorder="little", signed=False)) + h.update(wav_f32.tobytes(order="C")) + return h.hexdigest() + + def _get_ref_audio_artifacts(self, cache_key: str) -> dict[str, torch.Tensor | bool] | None: + entry = self._ref_audio_artifact_cache.get(cache_key) + if entry is None: + return None + # Touch for LRU. + self._ref_audio_artifact_cache.move_to_end(cache_key, last=True) + return entry + + def _put_ref_audio_artifacts( + self, + cache_key: str, + *, + ref_code: torch.Tensor | None = None, + ref_spk_embedding: torch.Tensor | None = None, + ) -> None: + if not cache_key: + return + entry = self._ref_audio_artifact_cache.get(cache_key) + if entry is None: + entry = {} + if isinstance(ref_code, torch.Tensor): + entry["ref_code"] = ref_code.detach().to("cpu", dtype=torch.long).contiguous() + if isinstance(ref_spk_embedding, torch.Tensor): + entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=torch.bfloat16).reshape(-1) + if not entry: + return + self._ref_audio_artifact_cache[cache_key] = entry + self._ref_audio_artifact_cache.move_to_end(cache_key, last=True) + while len(self._ref_audio_artifact_cache) > self._ref_audio_artifact_cache_max_entries: + self._ref_audio_artifact_cache.popitem(last=False) + @staticmethod def _first_value(value: object, default: object = None) -> object: if isinstance(value, list): @@ -1544,6 +1603,14 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # invalid ref_audio errors with its existing messages. continue info_dict[_NORMALIZED_REF_AUDIO_KEY] = (wav, sr) + cache_key = self._make_ref_audio_cache_key(wav, sr) + info_dict[_REF_AUDIO_CACHE_KEY] = cache_key + cached = self._get_ref_audio_artifacts(cache_key) + if cached is not None: + cached_ref_code = self._coerce_ref_code_tensor(cached.get("ref_code"), device=device) + if isinstance(cached_ref_code, torch.Tensor): + info_dict.setdefault("codes", {})[_PRECOMPUTED_REF_CODE_KEY] = cached_ref_code + continue groups.setdefault(int(sr), []).append((info_dict, wav, int(sr))) if pending_text or pending_ref_text: @@ -1587,6 +1654,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): continue info_dict.setdefault("codes", {})[_PRECOMPUTED_REF_CODE_KEY] = ref_code_t info_dict[_NORMALIZED_REF_AUDIO_KEY] = (wav, item_sr) + cache_key = info_dict.get(_REF_AUDIO_CACHE_KEY) + if isinstance(cache_key, str) and cache_key: + self._put_ref_audio_artifacts(cache_key, ref_code=ref_code_t) def _generate_icl_prompt( self, @@ -1773,6 +1843,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): voice_clone_prompt = _normalize_voice_clone_prompt(info_dict.get("voice_clone_prompt")) ref_audio_wav: np.ndarray | None = None ref_audio_sr: int | None = None + ref_audio_cache_key: str | None = None def _get_ref_audio() -> tuple[np.ndarray, int]: nonlocal ref_audio_wav, ref_audio_sr @@ -1792,6 +1863,32 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): ref_audio_wav, ref_audio_sr = self._normalize_ref_audio(ref_audio_list[0]) return ref_audio_wav, ref_audio_sr + def _get_ref_audio_cache_key() -> str: + nonlocal ref_audio_cache_key + if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key: + return ref_audio_cache_key + cached_key = info_dict.get(_REF_AUDIO_CACHE_KEY) + if isinstance(cached_key, str) and cached_key: + ref_audio_cache_key = cached_key + return ref_audio_cache_key + wav_np, sr = _get_ref_audio() + ref_audio_cache_key = self._make_ref_audio_cache_key(wav_np, sr) + info_dict[_REF_AUDIO_CACHE_KEY] = ref_audio_cache_key + return ref_audio_cache_key + + def _try_get_ref_audio_cache_key() -> str | None: + nonlocal ref_audio_cache_key + if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key: + return ref_audio_cache_key + cached_key = info_dict.get(_REF_AUDIO_CACHE_KEY) + if isinstance(cached_key, str) and cached_key: + ref_audio_cache_key = cached_key + return ref_audio_cache_key + ref_audio_list = info_dict.get("ref_audio") + if isinstance(ref_audio_list, list) and ref_audio_list: + return _get_ref_audio_cache_key() + return None + # Speaker cache: only for uploaded (named) speakers _speaker_cache_key = None if voice_clone_prompt is None: @@ -1849,12 +1946,24 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): if isinstance(ref_code_t, torch.Tensor): ref_code_len = int(ref_code_t.shape[0]) if ref_code_t is None and in_context_mode: - # Compute ref_code from ref_audio if not provided. - wav_np, sr = _get_ref_audio() - ref_code_t = self._encode_ref_audio_to_code(wav_np, sr).to(device=input_ids.device) - ref_code_len = int(ref_code_t.shape[0]) + cache_key = _get_ref_audio_cache_key() + cached = self._get_ref_audio_artifacts(cache_key) + cached_ref_code = None + if isinstance(cached, dict): + cached_ref_code = cached.get("ref_code") + ref_code_t = self._coerce_ref_code_tensor(cached_ref_code, device=input_ids.device) + if isinstance(ref_code_t, torch.Tensor): + ref_code_len = int(ref_code_t.shape[0]) + else: + wav_np, sr = _get_ref_audio() + ref_code_t = self._encode_ref_audio_to_code(wav_np, sr).to(device=input_ids.device) + ref_code_len = int(ref_code_t.shape[0]) + self._put_ref_audio_artifacts(cache_key, ref_code=ref_code_t) if isinstance(ref_code_t, torch.Tensor): ref_code_prompt = ref_code_t + cache_key = _try_get_ref_audio_cache_key() + if isinstance(cache_key, str) and cache_key: + self._put_ref_audio_artifacts(cache_key, ref_code=ref_code_prompt) # Speaker embedding: use prompt embed if provided; otherwise extract from audio. # NOTE: Do NOT use _as_singleton here — the embedding may be a plain @@ -1869,8 +1978,17 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # Plain list/array from API (survived msgspec IPC serialization). speaker_embed = torch.tensor(spk, dtype=torch.bfloat16, device=input_ids.device).view(1, 1, -1) else: - wav_np, sr = _get_ref_audio() - speaker_embed = self._extract_speaker_embedding(wav_np, sr).view(1, 1, -1) + cache_key = _get_ref_audio_cache_key() + cached = self._get_ref_audio_artifacts(cache_key) + cached_spk = None + if isinstance(cached, dict): + cached_spk = cached.get("ref_spk_embedding") + if isinstance(cached_spk, torch.Tensor): + speaker_embed = cached_spk.to(device=input_ids.device, dtype=torch.bfloat16).view(1, 1, -1) + else: + wav_np, sr = _get_ref_audio() + speaker_embed = self._extract_speaker_embedding(wav_np, sr).view(1, 1, -1) + self._put_ref_audio_artifacts(cache_key, ref_spk_embedding=speaker_embed) # Cache miss: store extraction result if _speaker_cache_key is not None and speaker_embed is not None: diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 8d058335e..3a8556273 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -155,6 +155,11 @@ class NPUARModelRunner(OmniNPUModelRunner): finally: set_cudagraph_capturing_enabled(False) + def _model_needs_full_prefix_hidden_states(self) -> bool: + """See gpu_ar_model_runner._model_needs_full_prefix_hidden_states.""" + model = getattr(self, "model", None) + return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True)) + def _maybe_update_prefix_cache( self, hidden_states: torch.Tensor, @@ -169,8 +174,9 @@ class NPUARModelRunner(OmniNPUModelRunner): type(multimodal_outputs), ) + hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None self.omni_prefix_cache.update_omni_tensor_prefix_cache( - hidden_states=hidden_states, + hidden_states=hs_for_cache, multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs, num_tokens_unpadded=num_tokens_unpadded, slot_mapping=self.input_batch.block_table[0].slot_mapping.cpu, @@ -185,12 +191,13 @@ class NPUARModelRunner(OmniNPUModelRunner): ) -> tuple[dict[str, torch.Tensor] | None, dict | None]: combined_hidden_states, combined_multimodal_outputs = None, None if self.omni_prefix_cache is not None: - combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states( - query_start_loc=self.query_start_loc.cpu, - input_batch=self.input_batch, - hidden_states=hidden_states, - num_scheduled_tokens=num_scheduled_tokens, - ) + if self._model_needs_full_prefix_hidden_states(): + combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states( + query_start_loc=self.query_start_loc.cpu, + input_batch=self.input_batch, + hidden_states=hidden_states, + num_scheduled_tokens=num_scheduled_tokens, + ) combined_multimodal_outputs = self.omni_prefix_cache.get_merged_multimodal_states( query_start_loc=self.query_start_loc.cpu, input_batch=self.input_batch, diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py index cf2837a7c..43e8fa110 100644 --- a/vllm_omni/utils/mm_outputs.py +++ b/vllm_omni/utils/mm_outputs.py @@ -68,12 +68,13 @@ def to_payload_element( if we need to avoid splitting nonempty lists prior to calling postprocess, which is the case for prefix cache. seq_len: Optional sequence length (i.e., dim 0 of hidden states). - This should be set to None in the prefix caching case, because - the condition that would be executed here is the same as the - criteria for being added to the multimodal outputs cache. + When set, a tensor whose first dimension equals seq_len is + sliced per request. The prefix cache passthrough also passes + the total scheduled token count here so 1D (seq_len,) metadata + that is intentionally not cached is still split per request. """ - # Prefix cache won't hit this case because this is the condition - # for being a mm_cache_key in the multimodal outputs tensor. + # Cached per-token tensors are merged elsewhere; here a first dim + # equal to seq_len means a per-request slice is required. if seq_len is not None and isinstance(element, torch.Tensor) and element.shape[0] == seq_len: return element[start:end].contiguous() # Every other case is shared between prefix cache (passthrough data) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 34d571c11..b3fa99fc5 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -203,6 +203,24 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): finally: set_cudagraph_capturing_enabled(False) + def _model_needs_full_prefix_hidden_states(self) -> bool: + """Opt-out hook for models whose postprocess only consumes the tail. + + When False, we skip both the per-step GPU->CPU hidden-state write into + OmniTensorPrefixCache and the merged-tensor reconstruction on hits; + postprocess receives the normal scheduled-token slice instead. Models + that need the full cached_prefix + new_tail span (default) are not + affected. + """ + model = getattr(self, "model", None) + return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True)) + + def _deferred_prefix_cache_mm_keys(self) -> set[str]: + """Model-declared multimodal keys whose prefix-cache writes are deferred.""" + model = getattr(self, "model", None) + keys = getattr(model, "deferred_prefix_cache_mm_keys", ()) + return set(keys or ()) + def _maybe_update_prefix_cache( self, hidden_states: torch.Tensor, @@ -225,12 +243,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): type(multimodal_outputs), ) + hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None self.omni_prefix_cache.update_omni_tensor_prefix_cache( - hidden_states=hidden_states, + hidden_states=hs_for_cache, multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs, num_tokens_unpadded=num_tokens_unpadded, slot_mapping=self.input_batch.block_table[0].slot_mapping.cpu, num_tokens_padded=num_tokens_padded, + skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(), ) def _maybe_get_combined_prefix_cache_tensors( @@ -246,12 +266,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # the prefix cached hidden states and multimodal states. combined_hidden_states, combined_multimodal_outputs = None, None if self.omni_prefix_cache is not None: - combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states( - query_start_loc=self.query_start_loc.cpu, - input_batch=self.input_batch, - hidden_states=hidden_states, - num_scheduled_tokens=num_scheduled_tokens, - ) + if ( + not self._model_needs_full_prefix_hidden_states() + and not self.omni_prefix_cache.has_prefix_cached_new_req_ids() + ): + return None, None + if self._model_needs_full_prefix_hidden_states(): + combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states( + query_start_loc=self.query_start_loc.cpu, + input_batch=self.input_batch, + hidden_states=hidden_states, + num_scheduled_tokens=num_scheduled_tokens, + ) combined_multimodal_outputs = self.omni_prefix_cache.get_merged_multimodal_states( query_start_loc=self.query_start_loc.cpu, input_batch=self.input_batch, @@ -313,6 +339,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if flush_ids: self.flush_full_payload_outputs(flush_ids) + if self.omni_prefix_cache is not None and scheduler_output.finished_req_ids: + self.omni_prefix_cache.commit_deferred_mm_outputs( + set(scheduler_output.finished_req_ids), + self.input_batch, + ) + if self.routed_experts_initialized: capturer = get_global_experts_capturer() if capturer is not None and hasattr(capturer, "finalize_pending_copy"): @@ -884,9 +916,22 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): dtype=np.int32, ) query_start_loc_cpu = self.query_start_loc.cpu + if self.omni_prefix_cache is not None: + deferred_mm_cache_keys = self._deferred_prefix_cache_mm_keys() + if deferred_mm_cache_keys: + self.omni_prefix_cache.stage_deferred_mm_outputs( + query_start_loc=query_start_loc_cpu, + input_batch=self.input_batch, + multimodal_outputs=( + flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs + ), + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + deferred_mm_cache_keys=deferred_mm_cache_keys, + ) pooler_output: list[dict[str, object]] | None = None if needs_pooler_payload: + mm_cpu = None # Prior to applying the post-processing func, extract # the prefix cached hidden states and multimodal states. if self.omni_prefix_cache is not None: @@ -899,8 +944,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): scheduler_output.num_scheduled_tokens, ) # Otherwise we don't have the mm CPU data yet, so we still need to build it - if self.omni_prefix_cache is None: - mm_cpu = build_mm_cpu(flatten_payload(multimodal_outputs)) + if self.omni_prefix_cache is None or combined_multimodal_outputs is None: + mm_cpu = build_mm_cpu(flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs) self._process_additional_information_updates( hidden_states, diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 5f76f7be1..9ba0b5f7b 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -332,6 +332,8 @@ class OmniGPUModelRunner(GPUModelRunner): self.requests.pop(req_id, None) self.model_intermediate_buffer.pop(req_id, None) self.num_prompt_logprobs.pop(req_id, None) + if self.omni_prefix_cache is not None: + self.omni_prefix_cache.discard_deferred_mm_outputs(req_id) if hasattr(self, "_downstream_payload_cache"): self._downstream_payload_cache.pop(req_id, None) if hasattr(self, "_talker_mtp_generators"):