From 5dd1d0a772179f2e36298091bbcd30f467329d4c Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Fri, 29 May 2026 09:51:54 +0800 Subject: [PATCH] [BugFix] Fix prefix-caching issue (#3726) Signed-off-by: amy-why-3459 --- tests/e2e/online_serving/test_qwen3_omni_expansion.py | 6 +++++- vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py | 9 +++++++-- vllm_omni/worker/gpu_ar_model_runner.py | 7 ++++++- 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 0e6cca111..aa176c077 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -61,7 +61,11 @@ test_params = [ model=model, stage_config_path=default_path, use_stage_cli=True, - server_args=["--no-async-chunk"], + server_args=[ + "--no-async-chunk", + "--stage-overrides", + '{"0": {"enable_prefix_caching": true}, "1": {"enable_prefix_caching": true}}', + ], ), id="default", ), diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index ba924a1ea..6cecfa587 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -673,8 +673,13 @@ class Qwen3OmniMoeForConditionalGeneration( span_len = input_ids.shape[0] update_dict: OmniPayload = {} - if span_len > 1: - # prefill + # Prefix caching can reduce a new request's remaining prefill span to a + # single token. Use the runner-provided phase flag instead of span_len. + is_prefill = bool(payload.get("_omni_is_prefill", span_len > 1)) + if is_prefill: + num_computed_tokens = payload.get("_omni_num_computed_tokens") + if num_computed_tokens is not None: + meta["num_processed_tokens"] = int(num_computed_tokens) input_ids, input_embeds, update_dict = self.talker_preprocess_prefill(input_ids, input_embeds, payload) code_predictor_codes = torch.zeros( (input_embeds.shape[0], self.talker.num_code_groups), diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 60c059c5a..f9ce18c6a 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -307,11 +307,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): ) hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None + # FIX: The .cpu attribute of slot_mapping is stale (not updated by the Triton + # _compute_slot_mapping_kernel which only writes to .gpu). We must use .gpu and + # sync back to CPU to get the correctly computed slot mapping. + slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu + slot_mapping_cpu = slot_mapping_gpu[:num_tokens_padded].cpu() self.omni_prefix_cache.update_omni_tensor_prefix_cache( 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, + slot_mapping=slot_mapping_cpu, num_tokens_padded=num_tokens_padded, skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(), )