diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index ed210dc9c..f1a3ce353 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -37,19 +37,6 @@ def get_chunk_config(config_path: str | None = None): return modify_stage_config(config_path, updates={"async_chunk": True}) -def get_prefix_caching_config(config_path: str): - """Create a stage config with prefix caching enabled on the thinker (stage 0).""" - path = modify_stage_config( - config_path, - updates={ - "stage_args": { - 0: {"engine_args.enable_prefix_caching": True}, - }, - }, - ) - return path - - # Platform-specific overrides live inside the new deploy yaml's ``platforms:`` # section, so a single ``_CI_DEPLOY`` path serves CUDA, ROCm, and XPU. # TODO: re-add VLLM_TEST_PD_MODE branch once the PD-disaggregation deploy @@ -59,22 +46,27 @@ if current_omni_platform.is_xpu(): stage_configs = [_CI_DEPLOY] else: # CUDA + ROCm MI325 share the same deploy config stage_configs = [get_chunk_config()] -prefix_caching_stage_configs = [get_prefix_caching_config(_CI_DEPLOY)] # Create parameter combinations for model and stage config test_params = [ OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs ] -# For prefix caching, we need to enable prompt token details so that we -# can determine if any tokens were cached. +# For prefix caching, we enable it on the thinker (stage 0) via CLI override +# and enable prompt token details so that we can determine if any tokens were cached. +BLOCK_SIZE = 16 prefix_test_params = [ OmniServerParams( model=model, - stage_config_path=stage_config, - server_args=["--enable-prompt-tokens-details"], # Enable prompt tokens details to get cached_tokens + stage_config_path=_CI_DEPLOY, + server_args=[ + "--block-size", + str(BLOCK_SIZE), + "--stage-overrides", + '{"0": {"enable_prefix_caching": true}}', + "--enable-prompt-tokens-details", + ], ) for model in models - for stage_config in prefix_caching_stage_configs ] @@ -180,16 +172,17 @@ def test_text_to_text_001(omni_server, openai_client) -> None: @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", prefix_test_params, indirect=True) -@pytest.mark.skip(reason="Flaky: identical greedy requests diverge on text_content (prefix caching path).") -def test_thinker_prefix_caching(omni_server, openai_client) -> None: +def test_thinker_prefix_caching(omni_server, openai_client, run_level) -> None: """ Test thinker prefix caching by sending identical requests with an image (i.e., a large shared prefix) and verifying that the second request uses cached tokens & produces the same output with greedy decoding. - NOTE: The seed for this test is used as a regression test for the issue linked below; - https://github.com/vllm-project/vllm-omni/issues/2833; without passing the sampling - params, this test will fail with the current default stage configs. + NOTE: The reason that we check against logprobs instead of direct text here is that + the outputs may still diverge a bit even though we set the seed and temperature. + This is mostly because the GEMM algorithm may vary based on the input tensors dims. + Because of this, we don't check the logprobs if it's a dummy load, since in that case + the top logprobs will all be very close. """ seed = 10 img_res = generate_synthetic_image(224, 224, seed=seed) @@ -200,19 +193,41 @@ def test_thinker_prefix_caching(omni_server, openai_client) -> None: content_text=get_prompt("text_image"), ) + top_k = 10 + sampling_params = {"seed": seed, "temperature": 0, "max_tokens": 8, "logprobs": top_k} request_config = { "model": omni_server.model, "messages": messages, "stream": False, "modalities": ["text"], - "sampling_params_list": [{"seed": seed, "temperature": 0, "max_tokens": 16}] * 3, + "logprobs": True, + "top_logprobs": top_k, + "sampling_params_list": [sampling_params] * 3, } - response_1 = openai_client.send_omni_request(request_config, request_num=1)[0] - response_2 = openai_client.send_omni_request(request_config, request_num=1)[0] + uncached_response = openai_client.send_omni_request(request_config, request_num=1)[0] + cached_response = openai_client.send_omni_request(request_config, request_num=1)[0] - # We should cache the vast majority of the prompt (image + up to last full block), - # and set seed + temperature, so the second request should give an identical - # response for the generated input image, even if we use dummy weights - assert response_2.cached_tokens is not None and response_2.cached_tokens > 0 - assert response_1.text_content == response_2.text_content + # Ensure that we have a prefix cache hit on the second request and that only the last + # partial block is uncached (since currently we don't cache partial blocks). + num_cached_tokens = cached_response.cached_tokens + num_prompt_tokens = cached_response.prompt_tokens + assert num_cached_tokens is not None and num_prompt_tokens is not None + num_uncached_tokens = num_prompt_tokens % BLOCK_SIZE + assert num_cached_tokens > 0 + assert num_cached_tokens % BLOCK_SIZE == 0 + assert (num_cached_tokens + num_uncached_tokens) == num_prompt_tokens + + # Ensure that we have logprobs and tokens were generated for both requests + assert uncached_response.logprobs is not None + assert cached_response.logprobs is not None + n_tokens = min(len(uncached_response.logprobs), len(cached_response.logprobs)) + assert n_tokens > 0 + + if run_level == "advanced_model": + # For each token index where both responses have an output, ensure that the greedy token + # predicted in the uncached case is in the top k logprobs for the cached case + for idx in range(n_tokens): + greedy_token = uncached_response.logprobs[idx].token + cached_top_k = {lp.token for lp in cached_response.logprobs[idx].top_logprobs} + assert greedy_token in cached_top_k diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 0b08e350b..45d86548f 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -615,7 +615,9 @@ class OmniResponse: e2e_latency: float | None = None success: bool = False error_message: str | None = None + prompt_tokens: int | None = None cached_tokens: int | None = None + logprobs: list | None = None @dataclass @@ -679,10 +681,12 @@ class OpenAIClientHandler: audio_data = choice.message.audio.data if hasattr(choice.message, "content") and choice.message.content is not None: text_content = choice.message.content - # Extract cached_tokens for prefix caching tests + # Extract cached & prompt token counts for prefix caching tests usage = getattr(chat_completion, "usage", None) - if usage and (details := getattr(usage, "prompt_tokens_details", None)): - result.cached_tokens = details.cached_tokens + if usage: + result.prompt_tokens = usage.prompt_tokens + if details := getattr(usage, "prompt_tokens_details", None): + result.cached_tokens = details.cached_tokens result.e2e_latency = time.perf_counter() - start_time audio_content = None if audio_data: @@ -690,6 +694,8 @@ class OpenAIClientHandler: audio_content = convert_audio_bytes_to_text(result.audio_bytes) result.text_content = text_content result.audio_content = audio_content + if chat_completion.choices and chat_completion.choices[0].logprobs is not None: + result.logprobs = chat_completion.choices[0].logprobs.content result.success = True except Exception as e: result.error_message = f"Non-stream processing error: {str(e)}" @@ -756,6 +762,10 @@ class OpenAIClientHandler: "stream": stream, "modalities": modalities, } + if "logprobs" in request_config: + create_kwargs["logprobs"] = request_config["logprobs"] + if "top_logprobs" in request_config: + create_kwargs["top_logprobs"] = request_config["top_logprobs"] if extra_body: create_kwargs["extra_body"] = extra_body