diff --git a/vllm_omni/core/prefix_cache.py b/vllm_omni/core/prefix_cache.py index 79f0990fd..25acfe8ed 100644 --- a/vllm_omni/core/prefix_cache.py +++ b/vllm_omni/core/prefix_cache.py @@ -391,27 +391,33 @@ class OmniTensorPrefixCache: self, query_start_loc: torch.Tensor, input_batch: InputBatch, - multimodal_outputs: dict, + multimodal_outputs: dict | None, num_scheduled_tokens: dict[str, int], ): """Get the merged multimodal states if hidden state prefix caching is enabled.""" combined_multimodal_outputs = {} + # Talkers that produce multimodal outputs only at decode (e.g. Higgs + # Audio v3's audio codes) have ``multimodal_outputs is None`` at the + # initial prefill step. Treat it as an empty mapping so the merger + # short-circuits cleanly and the cache merge still runs once codes + # start arriving. + mm_outputs = multimodal_outputs if multimodal_outputs is not None else {} # First get the prefix cached tensors that are present in the mm data for mm_key in self.mm_cache_keys: - if mm_key in multimodal_outputs: + if mm_key in mm_outputs: combined_multimodal_outputs[mm_key] = self._get_merged_tensors( query_start_loc=query_start_loc, input_batch=input_batch, cache=self.mm_outputs_cache[mm_key], - hidden_states=multimodal_outputs[mm_key], + hidden_states=mm_outputs[mm_key], num_scheduled_tokens=num_scheduled_tokens, ) # Then, get everything else (passthrough data); first, convert to CPU # tensors similarly to the non prefix cached path, and then populate # the subdicts mapping request IDs -> payload objects - passthrough_keys = set(multimodal_outputs.keys()) - self.mm_cache_keys - passthrough_mm_data = {k: v for k, v in multimodal_outputs.items() if k in passthrough_keys} + passthrough_keys = set(mm_outputs.keys()) - self.mm_cache_keys + passthrough_mm_data = {k: v for k, v in mm_outputs.items() if k in passthrough_keys} mm_cpu = build_mm_cpu(multimodal_outputs=passthrough_mm_data) for mm_key, mm_val in mm_cpu.items(): diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml index f6742f10a..c7016537b 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml +++ b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml @@ -33,7 +33,12 @@ stages: gpu_memory_utilization: 0.6 enforce_eager: true trust_remote_code: true - enable_prefix_caching: false + # Stage-0 prefix caching is on by default. The talker's HS opt-out and + # deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py) + # keep per-step bookkeeping negligible; any shared text prefix across + # concurrent requests is a net throughput win. Stage 1 (codec decoder) + # has no KV cache and stays disabled. + enable_prefix_caching: true async_scheduling: false max_num_batched_tokens: 4096 max_model_len: 8192