From 01f500a5bf84d84a11c41b76e42532a99f62c7fd Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Fri, 1 May 2026 22:47:24 +0800 Subject: [PATCH] [perf]Qwen3-Omni performance optimization (#3203) Signed-off-by: amy-why-3459 Co-authored-by: Hongsheng Liu --- tests/dfx/conftest.py | 23 +- tests/dfx/perf/scripts/run_benchmark.py | 8 +- tests/dfx/perf/tests/test_qwen_omni.json | 135 ++++++++++++ tests/dfx/stability/conftest.py | 10 +- tests/worker/test_gpu_ar_model_runner.py | 43 ++++ vllm_omni/benchmarks/metrics/metrics.py | 5 +- vllm_omni/core/sched/omni_ar_scheduler.py | 2 +- .../stage_input_processors/qwen3_omni.py | 81 ++++++-- vllm_omni/worker/gpu_ar_model_runner.py | 196 ++++++++++++------ vllm_omni/worker/gpu_model_runner.py | 6 + 10 files changed, 412 insertions(+), 97 deletions(-) create mode 100644 tests/worker/test_gpu_ar_model_runner.py diff --git a/tests/dfx/conftest.py b/tests/dfx/conftest.py index 29bb33302..0c59bed99 100644 --- a/tests/dfx/conftest.py +++ b/tests/dfx/conftest.py @@ -70,17 +70,17 @@ def _build_serve_args(serve_args: Any) -> list[str]: def create_unique_server_params( configs: list[dict[str, Any]], stage_configs_dir: Path, -) -> list[tuple[str, str, str | None, str | None, tuple[str, ...]]]: - """Return one row per unique server configuration (same 5-tuple shape as upstream). +) -> list[tuple[str, str, str | None, str | None, tuple[str, ...], bool]]: + """Return one row per unique server configuration. - ``(test_name, model, deploy_yaml_path, stage_overrides_json, extra_cli_args)``. + ``(test_name, model, deploy_yaml_path, stage_overrides_json, extra_cli_args, use_omni)``. JSON ``server_params.serve_args`` (dict/list) is expanded via ``_build_serve_args`` and **prepended** to ``extra_cli_args`` so perf / stability ``omni_server`` fixtures stay identical to main while still honoring ``serve_args`` in benchmark JSON. """ - unique_params: list[tuple[str, str, str | None, str | None, tuple[str, ...]]] = [] - seen: set[tuple[str, str, str | None, str | None, tuple[str, ...]]] = set() + unique_params: list[tuple[str, str, str | None, str | None, tuple[str, ...], bool]] = [] + seen: set[tuple[str, str, str | None, str | None, tuple[str, ...], bool]] = set() for config in configs: test_name = config["test_name"] server_params = config["server_params"] @@ -104,8 +104,16 @@ def create_unique_server_params( serve_flat = _build_serve_args(server_params.get("serve_args")) raw_extra = tuple(server_params.get("extra_cli_args") or ()) extra_cli_args = tuple(serve_flat) + raw_extra + use_omni = bool(server_params.get("use_omni", True)) - server_param = (test_name, model, stage_config_path, stage_overrides_json, extra_cli_args) + server_param = ( + test_name, + model, + stage_config_path, + stage_overrides_json, + extra_cli_args, + use_omni, + ) if server_param not in seen: seen.add(server_param) unique_params.append(server_param) @@ -145,8 +153,9 @@ def create_reliability_omni_server_params( model=model, stage_config_path=stage_config_path, server_args=server_args_by_name.get(test_name), + use_omni=use_omni, ) - for test_name, model, stage_config_path, _stage_overrides_json, _extra_cli_args in unique_params + for test_name, model, stage_config_path, _stage_overrides_json, _extra_cli_args, use_omni in unique_params ] diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index 9036508cb..c1f3264c1 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -65,11 +65,13 @@ def omni_server(request): Multi-stage initialization can take 10-20+ minutes. """ with _omni_server_lock: - test_name, model, stage_config_path, stage_overrides, extra_cli_args = request.param + test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = request.param print(f"Starting OmniServer with test: {test_name}, model: {model}") - server_args = ["--stage-init-timeout", "600", "--init-timeout", "900"] + server_args: list[str] = [] + if use_omni: + server_args += ["--stage-init-timeout", "600", "--init-timeout", "900"] # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py): # deploy-config sets the base; stage-overrides are applied on top. Both can be set. if stage_config_path: @@ -78,7 +80,7 @@ def omni_server(request): server_args = ["--stage-overrides", stage_overrides] + server_args if extra_cli_args: server_args = list(extra_cli_args) + server_args - with OmniServer(model, server_args) as server: + with OmniServer(model, server_args, use_omni=use_omni) as server: server.test_name = test_name print("OmniServer started successfully") yield server diff --git a/tests/dfx/perf/tests/test_qwen_omni.json b/tests/dfx/perf/tests/test_qwen_omni.json index eda9720c4..3c70e027c 100644 --- a/tests/dfx/perf/tests/test_qwen_omni.json +++ b/tests/dfx/perf/tests/test_qwen_omni.json @@ -311,5 +311,140 @@ } } ] + }, + { + "test_name": "test_qwen3_omni_vllm_text", + "server_params": { + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "use_omni": false, + "extra_cli_args": ["--no-enable-prefix-caching"] + }, + "benchmark_params": [ + { + "dataset_name": "random", + "backend": "vllm", + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "endpoint": "/v1/completions", + "num_prompts": [4, 16, 32, 64, 128], + "max_concurrency": [1, 4, 8, 16, 32], + "random_input_len": 2500, + "random_output_len": 900, + "temperature": 0.4, + "top_p": 0.9, + "top_k": 1, + "seed": 42, + "repetition_penalty": 1.05, + "ignore_eos": true, + "extra_body": { + "modalities": ["text"] + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000] + } + }, + { + "dataset_name": "random-mm", + "backend": "vllm", + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "endpoint": "/v1/completions", + "num_prompts": [10], + "request_rate": [0.1], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "temperature": 0.4, + "top_p": 0.9, + "top_k": 1, + "seed": 42, + "repetition_penalty": 1.05, + "ignore_eos": true, + "extra_body": { + "modalities": ["text"] + }, + "random_mm_base_items_per_request": 1, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "audio": 1 + }, + "random_mm_bucket_config": { + "(0, 60, 3)": 1.0 + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [2000] + } + }, + { + "dataset_name": "random-mm", + "backend": "vllm", + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "endpoint": "/v1/completions", + "num_prompts": [40], + "request_rate": [0.5], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "temperature": 0.4, + "top_p": 0.9, + "top_k": 1, + "seed": 42, + "repetition_penalty": 1.05, + "ignore_eos": true, + "extra_body": { + "modalities": ["text"] + }, + "random_mm_base_items_per_request": 2, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.5, + "(720, 1280, 2)": 0.5 + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [6000] + } + }, + { + "dataset_name": "random-mm", + "backend": "vllm", + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "endpoint": "/v1/completions", + "num_prompts": [100], + "request_rate": [1.0], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "temperature": 0.4, + "top_p": 0.9, + "top_k": 1, + "seed": 42, + "repetition_penalty": 1.05, + "ignore_eos": true, + "extra_body": { + "modalities": ["text"] + }, + "random_mm_base_items_per_request": 3, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1, + "audio": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.34, + "(720, 1280, 2)": 0.33, + "(0, 60, 3)": 0.33 + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [6000] + } + } + ] } ] diff --git a/tests/dfx/stability/conftest.py b/tests/dfx/stability/conftest.py index 30718d4bf..eaa755d45 100644 --- a/tests/dfx/stability/conftest.py +++ b/tests/dfx/stability/conftest.py @@ -35,20 +35,22 @@ def omni_server(request: pytest.FixtureRequest): """Start OmniServer for stability tests, with per-module timeout override.""" timeout_args = getattr(request.module, "STABILITY_SERVER_TIMEOUT_ARGS", DEFAULT_STABILITY_SERVER_TIMEOUT_ARGS) with _omni_server_lock: - # Same 5-tuple and CLI composition as ``tests/dfx/perf/scripts/run_benchmark.py`` on main; + # Same tuple and CLI composition as ``tests/dfx/perf/scripts/run_benchmark.py``; # ``serve_args`` from JSON are folded into ``extra_cli_args`` inside # ``create_unique_server_params``. - test_name, model, deploy_path, stage_overrides, extra_cli_args = request.param + test_name, model, deploy_path, stage_overrides, extra_cli_args, use_omni = request.param print(f"Starting OmniServer with test: {test_name}, model: {model}") - server_args = list(timeout_args) + server_args: list[str] = [] + if use_omni: + server_args += list(timeout_args) if deploy_path: server_args = ["--deploy-config", deploy_path] + server_args if stage_overrides: server_args = ["--stage-overrides", stage_overrides] + server_args if extra_cli_args: server_args = list(extra_cli_args) + server_args - with OmniServer(model, server_args) as server: + with OmniServer(model, server_args, use_omni=use_omni) as server: server.test_name = test_name print("OmniServer started successfully") yield server diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py new file mode 100644 index 000000000..87af8982e --- /dev/null +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -0,0 +1,43 @@ +from types import SimpleNamespace + +import pytest + +from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_runner(engine_output_type: str | None, downstream_req_ids: set[str]) -> GPUARModelRunner: + runner = object.__new__(GPUARModelRunner) + runner.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(engine_output_type=engine_output_type), + ) + runner._request_needs_downstream_stage_payload = lambda rid: rid in downstream_req_ids + return runner + + +def test_resolve_pooler_payload_req_ids_audio_terminal_stage_keeps_payload(): + runner = _make_runner(engine_output_type="audio", downstream_req_ids=set()) + + engine_output_type, payload_req_ids = GPUARModelRunner._resolve_pooler_payload_req_ids(runner, ["r1", "r2"]) + + assert engine_output_type == "audio" + assert payload_req_ids == ["r1", "r2"] + + +def test_resolve_pooler_payload_req_ids_text_terminal_stage_drops_payload(): + runner = _make_runner(engine_output_type="text", downstream_req_ids=set()) + + engine_output_type, payload_req_ids = GPUARModelRunner._resolve_pooler_payload_req_ids(runner, ["r1", "r2"]) + + assert engine_output_type == "text" + assert payload_req_ids == [] + + +def test_resolve_pooler_payload_req_ids_downstream_stage_uses_filtered_requests(): + runner = _make_runner(engine_output_type="latent", downstream_req_ids={"r2"}) + + engine_output_type, payload_req_ids = GPUARModelRunner._resolve_pooler_payload_req_ids(runner, ["r1", "r2", "r3"]) + + assert engine_output_type == "latent" + assert payload_req_ids == ["r2"] diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py index dbf764698..f320fffe9 100644 --- a/vllm_omni/benchmarks/metrics/metrics.py +++ b/vllm_omni/benchmarks/metrics/metrics.py @@ -188,7 +188,10 @@ def calculate_metrics( total_input += outputs[i].prompt_len tpot = 0 if output_len > 1: - latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft + try: + latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft + except Exception: + latency_minus_ttft = outputs[i].latency - outputs[i].ttft tpot = latency_minus_ttft / (output_len - 1) tpots.append(tpot) # Note: if output_len <= 1, we regard tpot as 0 for goodput diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 9dfbf5dff..677d8f671 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -9,8 +9,8 @@ from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.distributed.kv_events import KVEventBatch from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.logger import init_logger +from vllm.v1.core.sched.async_scheduler import AsyncScheduler as VLLMScheduler from vllm.v1.core.sched.output import SchedulerOutput -from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.core.sched.utils import remove_all from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs from vllm.v1.metrics.perf import PerfStats 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 8018df343..f7da44067 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -28,6 +28,17 @@ _EMBED_LAYER_KEY = "0" _HIDDEN_LAYER_KEY = "24" +def _layer_tensor(layers: dict[Any, Any], key: str) -> torch.Tensor | None: + """Fetch layer tensor with tolerant key lookup (str/int).""" + if not isinstance(layers, dict): + return None + key_int = int(key) + val = layers.get(key_int) + if val is None: + val = layers.get(key) + return val if isinstance(val, torch.Tensor) else None + + def _compute_talker_prompt_ids_length(info: OmniPayload, device: torch.device | str = "cuda") -> int: im_start_token_id = 151644 system_token_id = 8948 @@ -300,9 +311,23 @@ def thinker2talker_async_chunk( request_id = request.external_req_id chunk_id = transfer_manager.put_req_chunk[request_id] + if not isinstance(pooling_output, dict): + logger.debug("thinker2talker_async_chunk: skip non-dict pooling_output for req=%s", request_id) + return None + thinker_hs = pooling_output.get("hidden_states", {}) - thinker_layers = thinker_hs.get("layers", {}) - thinker_embed = pooling_output.get("embed", {}) + thinker_layers = thinker_hs.get("layers", {}) if isinstance(thinker_hs, dict) else {} + thinker_embed = pooling_output.get("embed", {}) if isinstance(pooling_output.get("embed", {}), dict) else {} + thinker_emb = _layer_tensor(thinker_layers, _EMBED_LAYER_KEY) + thinker_hid = _layer_tensor(thinker_layers, _HIDDEN_LAYER_KEY) + if thinker_emb is None or thinker_hid is None: + logger.debug( + "thinker2talker_async_chunk: missing thinker layers for req=%s (embed=%s hidden=%s)", + request_id, + thinker_emb is not None, + thinker_hid is not None, + ) + return None if chunk_id == 0: all_token_ids = request.all_token_ids # prefill + decode @@ -312,13 +337,19 @@ def thinker2talker_async_chunk( prompt_token_ids = _ensure_list(prompt_token_ids) payload: OmniPayload = { "embed": { - "prefill": thinker_layers[int(_EMBED_LAYER_KEY)].detach().cpu(), + "prefill": thinker_emb.detach().cpu(), # Provide thinker-side TTS token embeddings for talker projection - "tts_bos": thinker_embed["tts_bos"].detach().cpu(), - "tts_eos": thinker_embed["tts_eos"].detach().cpu(), - "tts_pad": thinker_embed["tts_pad"].detach().cpu(), + "tts_bos": thinker_embed.get("tts_bos").detach().cpu() + if isinstance(thinker_embed.get("tts_bos"), torch.Tensor) + else None, + "tts_eos": thinker_embed.get("tts_eos").detach().cpu() + if isinstance(thinker_embed.get("tts_eos"), torch.Tensor) + else None, + "tts_pad": thinker_embed.get("tts_pad").detach().cpu() + if isinstance(thinker_embed.get("tts_pad"), torch.Tensor) + else None, }, - "hidden_states": {"output": thinker_layers[int(_HIDDEN_LAYER_KEY)].detach().cpu()}, + "hidden_states": {"output": thinker_hid.detach().cpu()}, "ids": {"all": all_token_ids, "prompt": prompt_token_ids}, "meta": {"finished": torch.tensor(is_finished, dtype=torch.bool)}, } @@ -366,12 +397,12 @@ def thinker2talker_async_chunk( if output_token_ids: talker_additional_info["meta"]["override_keys"] = [("embed", "decode"), ("ids", "output")] - talker_additional_info["embed"] = {"decode": thinker_layers[int(_EMBED_LAYER_KEY)].detach().cpu()} + talker_additional_info["embed"] = {"decode": thinker_emb.detach().cpu()} talker_additional_info["ids"] = {"output": output_token_ids} else: # When prefilling a chunked thinker, thinker_hidden_states needs to be updated. - talker_additional_info["embed"] = {"prefill": thinker_layers[0].detach().cpu()} - talker_additional_info["hidden_states"] = {"output": thinker_layers[24].detach().cpu()} + talker_additional_info["embed"] = {"prefill": thinker_emb.detach().cpu()} + talker_additional_info["hidden_states"] = {"output": thinker_hid.detach().cpu()} return talker_additional_info @@ -431,11 +462,20 @@ def thinker2talker( thinker_sequences = prompt_token_ids + output_ids thinker_input_ids = prompt_token_ids new_seq_length = len(prompt_token_ids + output_ids) - 1 - thinker_mm: OmniPayload = output.multimodal_output + thinker_mm_raw = getattr(output, "multimodal_output", None) + if not isinstance(thinker_mm_raw, dict): + logger.debug("thinker2talker: skip req=%s due to empty multimodal_output", req_id) + continue + thinker_mm: OmniPayload = thinker_mm_raw mm_hs = thinker_mm.get("hidden_states", {}) - mm_layers = mm_hs.get("layers", {}) - thinker_emb = mm_layers[int(_EMBED_LAYER_KEY)].detach().to(device=device, dtype=torch.float)[-new_seq_length:] - thinker_hid = mm_layers[int(_HIDDEN_LAYER_KEY)].detach().to(device=device, dtype=torch.float)[-new_seq_length:] + mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, dict) else {} + emb_layer = _layer_tensor(mm_layers, _EMBED_LAYER_KEY) + hid_layer = _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY) + if emb_layer is None or hid_layer is None: + logger.debug("thinker2talker: skip req=%s due to missing hidden-state layers", req_id) + continue + thinker_emb = emb_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:] + thinker_hid = hid_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:] prefill_mm: dict[str, Any] | None = None if prefill_stage is not None: @@ -507,7 +547,11 @@ def talker2code2wav_async_chunk( """ Pooling version. """ + if not isinstance(pooling_output, dict): + return None talker_codes = pooling_output.get("codes", {}) + if not isinstance(talker_codes, dict): + return None code_predictor_codes = talker_codes.get("audio") if code_predictor_codes is None: return None @@ -600,7 +644,14 @@ def talker2code2wav( is_streaming_session = bool(getattr(streaming_context, "enabled", False)) if is_streaming_session: seq_len = _get_streaming_codec_delta_len(cur_seq_len, req_id, talker_output, streaming_context) - mm: OmniPayload = output.multimodal_output + mm_raw = getattr(output, "multimodal_output", None) + if not isinstance(mm_raw, dict): + logger.debug("talker2code2wav: skip req=%s due to empty multimodal_output", req_id) + continue + mm: OmniPayload = mm_raw + if "codes" not in mm or not isinstance(mm.get("codes"), dict) or "audio" not in mm["codes"]: + logger.debug("talker2code2wav: skip req=%s due to missing codes.audio", req_id) + continue # Extract codec codes from talker output # Expected shape: [8, seq_len] (8-layer RVQ codes) codec_codes = ( diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index fd3af5629..93c38b9b1 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -81,6 +81,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False) # Initialize KV cache manager (preserve vllm_config fallback behavior) self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config) + self._downstream_payload_cache: dict[str, bool] = {} def _make_buffer(self, *size, dtype, numpy=True): # Prevent ray from pinning the buffer due to large size @@ -142,6 +143,38 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): return sampling_metadata return replace(sampling_metadata, output_token_ids=output_token_ids) + def _request_final_stage_id(self, req_id: str) -> int | None: + info = self.model_intermediate_buffer.get(req_id) + if not isinstance(info, dict): + req_state = self.requests.get(req_id) + info = getattr(req_state, "additional_information_cpu", None) + if not isinstance(info, dict): + return None + val = info.get("omni_final_stage_id") + try: + return int(val) + except (TypeError, ValueError): + return None + + def _request_needs_downstream_stage_payload(self, req_id: str) -> bool: + cached = self._downstream_payload_cache.get(req_id) + if cached is not None: + return cached + # Conservative default: keep payload if marker is missing. + final_stage_id = self._request_final_stage_id(req_id) + needs_payload = final_stage_id is None or final_stage_id > 0 + self._downstream_payload_cache[req_id] = needs_payload + return needs_payload + + def _resolve_pooler_payload_req_ids(self, req_ids_output_copy: list[str]) -> tuple[str, list[str]]: + downstream_req_ids = [rid for rid in req_ids_output_copy if self._request_needs_downstream_stage_payload(rid)] + engine_output_type = (self.vllm_config.model_config.engine_output_type or "").lower() + # Single-stage AR TTS models (e.g. VoxCPM2) finish on this stage but still + # need multimodal payloads for final audio postprocess/output. + if engine_output_type == "audio" and not downstream_req_ids: + downstream_req_ids = req_ids_output_copy + return engine_output_type, downstream_req_ids + def capture_model(self) -> int: result = super().capture_model() self._capture_talker_mtp_graphs() @@ -833,7 +866,20 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): kv_connector_output = self.kv_connector_output self.kv_connector_output = None - hidden_states_cpu = hidden_states.detach().to("cpu").contiguous() + engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy) + needs_pooler_payload = len(downstream_req_ids) > 0 + downstream_req_id_set = set(downstream_req_ids) + hidden_states_cpu = None + req_hidden_states_cpu: dict[str, torch.Tensor] | None = None + if needs_pooler_payload: + num_valid_tokens = min( + int(scheduler_output.total_num_scheduled_tokens), + int(hidden_states.shape[0]), + ) + if len(downstream_req_ids) == len(req_ids_output_copy): + hidden_states_cpu = hidden_states[:num_valid_tokens].detach().to("cpu").contiguous() + else: + req_hidden_states_cpu = {} num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None) if num_scheduled_tokens_np is None: req_ids = self.input_batch.req_ids @@ -841,78 +887,96 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids], dtype=np.int32, ) + query_start_loc_cpu = self.query_start_loc.cpu - # Prior to applying the post-processing func, extract - # the prefix cached hidden states and multimodal states. - if self.omni_prefix_cache is not None: - ( - combined_hidden_states, - combined_multimodal_outputs, - ) = self._maybe_get_combined_prefix_cache_tensors( + pooler_output: list[dict[str, object]] | None = None + if needs_pooler_payload: + # Prior to applying the post-processing func, extract + # the prefix cached hidden states and multimodal states. + if self.omni_prefix_cache is not None: + ( + combined_hidden_states, + combined_multimodal_outputs, + ) = self._maybe_get_combined_prefix_cache_tensors( + hidden_states, + multimodal_outputs, + 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)) + + self._process_additional_information_updates( hidden_states, multimodal_outputs, - 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)) - - self._process_additional_information_updates( - hidden_states, - multimodal_outputs, - num_scheduled_tokens_np, - scheduler_output, - combined_hidden_states, - combined_multimodal_outputs, - ) - - pooler_output: list[dict[str, object]] = [] - for rid in req_ids_output_copy: - idx = req_id_to_index_output_copy[rid] - start = int(self.query_start_loc.cpu[idx]) - sched = int(num_scheduled_tokens_np[idx]) - end = start + sched - # If prefix cache is enabled, we have already split everything - # by request and converted the states to CPU tensors - req_hidden_states = self._resolve_req_hidden_states( - hidden_states_cpu, + num_scheduled_tokens_np, + scheduler_output, combined_hidden_states, - rid, - start, - end, + combined_multimodal_outputs, + req_ids_filter=downstream_req_id_set, ) - payload: dict[str, object] = {"hidden": req_hidden_states} - mm_payload: dict[str, object] = {} - if combined_multimodal_outputs or mm_cpu: - if combined_multimodal_outputs: - # Prefix cache enabled; all items have already been processed - # and split apart for each request as needed, and all tensors - # have already been detached to the CPU. The only exception is - # lists, which we keep as passthrough data for consistent behavior - # in postprocess. - for mm_key in combined_multimodal_outputs.keys(): - value = combined_multimodal_outputs[mm_key][rid] - if isinstance(value, list): - mm_payload[mm_key] = value[idx] if idx < len(value) else value[0] - else: - mm_payload[mm_key] = value + if req_hidden_states_cpu is not None and combined_hidden_states is None: + for rid in downstream_req_ids: + idx = req_id_to_index_output_copy[rid] + start = int(query_start_loc_cpu[idx]) + sched = int(num_scheduled_tokens_np[idx]) + end = start + sched + req_hidden_states_cpu[rid] = hidden_states[start:end].detach().to("cpu").contiguous() + pooler_output = [] + for rid in req_ids_output_copy: + if rid not in downstream_req_id_set: + pooler_output.append({}) + continue + idx = req_id_to_index_output_copy[rid] + start = int(query_start_loc_cpu[idx]) + sched = int(num_scheduled_tokens_np[idx]) + end = start + sched + # If prefix cache is enabled, we have already split everything + # by request and converted the states to CPU tensors + if req_hidden_states_cpu is not None and combined_hidden_states is None: + req_hidden_states = req_hidden_states_cpu[rid] else: - # Prefix cache disabled; we still need to process the data - for mm_key, mm_val in mm_cpu.items(): - mm_payload[mm_key] = to_payload_element( - element=mm_val, - idx=idx, - start=start, - end=end, - pass_lists_through=False, - seq_len=seq_len, - ) - payload.update(mm_payload) - # Flatten nested dicts to dotted keys so pooling_output - # stays dict[str, torch.Tensor] for msgspec serialization. - pooler_output.append(flatten_payload(payload)) + req_hidden_states = self._resolve_req_hidden_states( + hidden_states_cpu, + combined_hidden_states, + rid, + start, + end, + ) + payload: dict[str, object] = {"hidden": req_hidden_states} + + mm_payload: dict[str, object] = {} + if combined_multimodal_outputs or mm_cpu: + if combined_multimodal_outputs: + # Prefix cache enabled; all items have already been processed + # and split apart for each request as needed, and all tensors + # have already been detached to the CPU. The only exception is + # lists, which we keep as passthrough data for consistent behavior + # in postprocess. + for mm_key in combined_multimodal_outputs.keys(): + value = combined_multimodal_outputs[mm_key][rid] + if isinstance(value, list): + mm_payload[mm_key] = value[idx] if idx < len(value) else value[0] + else: + mm_payload[mm_key] = value + + else: + # Prefix cache disabled; we still need to process the data + for mm_key, mm_val in mm_cpu.items(): + mm_payload[mm_key] = to_payload_element( + element=mm_val, + idx=idx, + start=start, + end=end, + pass_lists_through=False, + seq_len=seq_len, + ) + payload.update(mm_payload) + # Flatten nested dicts to dotted keys so pooling_output + # stays dict[str, torch.Tensor] for msgspec serialization. + pooler_output.append(flatten_payload(payload)) with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): if self.routed_experts_initialized: capturer = RoutedExpertsCapturer.get_instance() @@ -926,7 +990,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): sampled_token_ids=valid_sampled_token_ids, logprobs=logprobs_lists, prompt_logprobs_dict=prompt_logprobs_dict, - pooler_output=(pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None), + pooler_output=(pooler_output if engine_output_type != "text" and needs_pooler_payload else None), kv_connector_output=kv_connector_output, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, num_nans_in_logits=num_nans_in_logits, diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 13d9c66dc..0f7ed429e 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -268,6 +268,9 @@ 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 hasattr(self, "_downstream_payload_cache"): + self._downstream_payload_cache.pop(req_id, None) + if hasattr(self, "late_interaction_runner"): self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids) # Remove the finished requests from the persistent batch. @@ -1060,6 +1063,7 @@ class OmniGPUModelRunner(GPUModelRunner): scheduler_output: "SchedulerOutput", combined_hidden_states: dict[str, torch.Tensor] | None = None, combined_multimodal_outputs: dict[str, object] | None = None, + req_ids_filter: set[str] | None = None, ) -> None: """Process model-provided per-request updates and merge into model_intermediate_buffer.""" try: @@ -1067,6 +1071,8 @@ class OmniGPUModelRunner(GPUModelRunner): # TODO(Peiqi): do we have a more elegant way to do this? if hasattr(self.model, "has_postprocess") and self.model.has_postprocess: for req_index, req_id in enumerate(self.input_batch.req_ids): + if req_ids_filter is not None and req_id not in req_ids_filter: + continue req_infos = self.model_intermediate_buffer.get(req_id, {}) if combined_hidden_states: # Combined hidden states contains all hidden states for every request