From 89b527e53301878bba33ba4a5a30a1c3207376dc Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Mon, 9 Mar 2026 17:29:38 +0800 Subject: [PATCH] Dev/rebase v0170 (#1639) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Taichang Zhou Signed-off-by: tzhouam Signed-off-by: princepride Signed-off-by: Zhou Taichang Co-authored-by: 汪志鹏 Co-authored-by: Hongsheng Liu Co-authored-by: Cyrus Leung --- docker/Dockerfile.ci | 3 +- docker/Dockerfile.rocm | 2 +- .../diffusion/attention/test_attention_sp.py | 3 +- .../omni_connectors/test_basic_connectors.py | 9 +- tests/e2e/online_serving/test_qwen3_tts.py | 4 +- tests/e2e/stage_configs/qwen2_5_omni_ci.yaml | 5 +- tests/e2e/stage_configs/qwen3_omni_ci.yaml | 2 +- tests/entrypoints/test_omni_llm.py | 70 +++++ tests/perf/stage_configs/qwen3_omni.yaml | 2 +- .../test_gpu_generation_model_runner.py | 1 + vllm_omni/benchmarks/metrics/metrics.py | 22 +- vllm_omni/benchmarks/patch/patch.py | 7 +- vllm_omni/config/model.py | 2 + vllm_omni/core/sched/omni_ar_scheduler.py | 7 +- .../core/sched/omni_generation_scheduler.py | 16 +- vllm_omni/core/sched/output.py | 16 +- vllm_omni/diffusion/utils/hf_utils.py | 4 +- .../distributed/omni_connectors/adapter.py | 16 +- .../connectors/shm_connector.py | 1 + .../chunk_transfer_adapter.py | 2 +- .../omni_connectors/utils/serialization.py | 17 +- vllm_omni/engine/__init__.py | 14 +- vllm_omni/engine/arg_utils.py | 4 + vllm_omni/engine/input_processor.py | 135 +++++---- vllm_omni/engine/output_processor.py | 15 +- vllm_omni/entrypoints/async_omni.py | 29 +- vllm_omni/entrypoints/async_omni_llm.py | 12 +- vllm_omni/entrypoints/omni_llm.py | 48 +++- vllm_omni/entrypoints/omni_stage.py | 16 +- vllm_omni/entrypoints/openai/api_server.py | 1 + vllm_omni/entrypoints/openai/serving_chat.py | 140 +++++++-- vllm_omni/inputs/preprocess.py | 46 ++- .../qwen3_omni/qwen3_omni_moe_talker.py | 2 +- .../qwen3_omni/qwen3_omni_moe_thinker.py | 130 ++++++++- vllm_omni/request.py | 1 - vllm_omni/worker/gpu_ar_model_runner.py | 36 ++- .../worker/gpu_generation_model_runner.py | 86 ++++-- vllm_omni/worker/gpu_model_runner.py | 267 ++++++++++-------- 38 files changed, 828 insertions(+), 365 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 4ce8af017..1a6889663 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,9 +1,8 @@ ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.16.0 +ARG VLLM_BASE_TAG=v0.17.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} - COPY . . # Install system dependencies diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 390bae400..e9db9ee49 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.16.0 +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0 FROM ${BASE_IMAGE} AS base # Declare a variable to know if we want to use the nightly build or the stable build. diff --git a/tests/diffusion/attention/test_attention_sp.py b/tests/diffusion/attention/test_attention_sp.py index 447d1ea93..234ac1563 100644 --- a/tests/diffusion/attention/test_attention_sp.py +++ b/tests/diffusion/attention/test_attention_sp.py @@ -414,7 +414,8 @@ def ulysses_attention_on_test_model( # Generate and save full input data with fixed seed # Reinitialize RNG to ensure reproducibility torch.manual_seed(42) - current_omni_platform.seed_everything(42) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(42) full_hidden_states = torch.randn( (batch_size, seq_len, hidden_size), dtype=dtype, diff --git a/tests/distributed/omni_connectors/test_basic_connectors.py b/tests/distributed/omni_connectors/test_basic_connectors.py index 7b416be72..1b1965355 100644 --- a/tests/distributed/omni_connectors/test_basic_connectors.py +++ b/tests/distributed/omni_connectors/test_basic_connectors.py @@ -66,14 +66,13 @@ def shm_connector(): def test_put_get_inline(shm_connector): - """Test inline transfer for small data.""" + """Test transfer for small data (currently always uses SHM).""" data = {"small": "data"} - # Ensure data is smaller than threshold (100 bytes) success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data) assert success is True - assert "inline_bytes" in metadata - assert "shm" not in metadata + assert "shm" in metadata + assert "size" in metadata # Retrieve retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata) @@ -87,7 +86,7 @@ def test_put_get_shm(mocker: MockerFixture, shm_connector, monkeypatch: pytest.M data = {"large": "x" * 200} # Mock SHM return values - mock_handle = {"name": "test_shm", "size": 200} + mock_handle = {"name": "req_2", "size": 200} mock_write = mocker.MagicMock(return_value=mock_handle) monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_write_bytes", mock_write) diff --git a/tests/e2e/online_serving/test_qwen3_tts.py b/tests/e2e/online_serving/test_qwen3_tts.py index 2ab5cdd4b..fd0ef7662 100644 --- a/tests/e2e/online_serving/test_qwen3_tts.py +++ b/tests/e2e/online_serving/test_qwen3_tts.py @@ -28,7 +28,7 @@ def get_stage_config(name: str = "qwen3_tts.yaml"): return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name) -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def omni_server(): """Start vLLM-Omni server with CustomVoice model.""" stage_config_path = get_stage_config() @@ -225,7 +225,7 @@ class TestQwen3TTSAPIEndpoints: assert len(data["data"]) > 0 -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def omni_server_no_async_chunk(): """Start vLLM-Omni server with non-async-chunk config.""" stage_config_path = get_stage_config("qwen3_tts_no_async_chunk.yaml") diff --git a/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml b/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml index 59bdc9df9..c11de3e5d 100644 --- a/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml +++ b/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml @@ -74,12 +74,13 @@ stage_args: model_arch: Qwen2_5OmniForConditionalGeneration worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - gpu_memory_utilization: 0.15 + gpu_memory_utilization: 0.9 #increase the gpu memory utilization to enable the test on H800 enforce_eager: true trust_remote_code: true enable_prefix_caching: false engine_output_type: audio - max_num_batched_tokens: 100000 + max_num_batched_tokens: 4096 + max_model_len: 4096 engine_input_source: [1] final_output: true final_output_type: audio diff --git a/tests/e2e/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/qwen3_omni_ci.yaml index a557dffd8..6f1293f96 100644 --- a/tests/e2e/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/stage_configs/qwen3_omni_ci.yaml @@ -48,7 +48,7 @@ stage_args: model_arch: Qwen3OmniMoeForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.6 + gpu_memory_utilization: 0.5 enforce_eager: false trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py index 4ec98be90..b11b14c89 100644 --- a/tests/entrypoints/test_omni_llm.py +++ b/tests/entrypoints/test_omni_llm.py @@ -1,3 +1,4 @@ +import inspect import uuid import warnings from queue import Empty, Queue @@ -1194,3 +1195,72 @@ def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: Mo # Verify stop_stage_worker was called (process should be set) assert omni.stage_list[0]._proc is not None + + +# --------------------------------------------------------------------------- +# Signature compatibility tests — catch upstream API drift early +# --------------------------------------------------------------------------- + +_INHERITANCE_PAIRS: list[tuple[str, str, str, str]] = [ + ( + "vllm.entrypoints.llm", + "LLM", + "vllm_omni.entrypoints.omni_llm", + "OmniLLM", + ), +] + + +def _import_class(module_path: str, class_name: str): + mod = __import__(module_path, fromlist=[class_name]) + return getattr(mod, class_name) + + +@pytest.mark.parametrize( + "up_mod,up_cls,omni_mod,omni_cls", + _INHERITANCE_PAIRS, + ids=[f"{pair[3]}({pair[1]})" for pair in _INHERITANCE_PAIRS], +) +def test_overridden_method_signatures_compatible(up_mod: str, up_cls: str, omni_mod: str, omni_cls: str): + """All params accepted by a base-class method must also be accepted by + the overriding method in the vllm-omni subclass. This prevents + TypeError at runtime when upstream callers pass new arguments.""" + BaseCls = _import_class(up_mod, up_cls) + OmniCls = _import_class(omni_mod, omni_cls) + + failures: list[str] = [] + for name in dir(OmniCls): + if name.startswith("__") and name != "__init__": + continue + omni_method = getattr(OmniCls, name, None) + base_method = getattr(BaseCls, name, None) + if not (callable(omni_method) and callable(base_method)): + continue + if omni_method is base_method: + continue + try: + base_sig = inspect.signature(base_method) + omni_sig = inspect.signature(omni_method) + except (ValueError, TypeError): + continue + + omni_has_var_keyword = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in omni_sig.parameters.values()) + + base_params = base_sig.parameters + omni_params = set(omni_sig.parameters.keys()) + + missing = [] + for pname, param in base_params.items(): + if pname in omni_params: + continue + if omni_has_var_keyword and param.kind in ( + inspect.Parameter.KEYWORD_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ): + continue + missing.append(pname) + + if missing: + failures.append(f"{omni_cls}.{name}() missing params {sorted(missing)}; base={base_sig}, omni={omni_sig}") + + assert not failures, f"Signature mismatches found ({len(failures)}):\n" + "\n".join(f" - {f}" for f in failures) diff --git a/tests/perf/stage_configs/qwen3_omni.yaml b/tests/perf/stage_configs/qwen3_omni.yaml index 802f8dd24..01ca7d9fb 100644 --- a/tests/perf/stage_configs/qwen3_omni.yaml +++ b/tests/perf/stage_configs/qwen3_omni.yaml @@ -85,7 +85,7 @@ stage_args: engine_output_type: audio # Final output: audio waveform gpu_memory_utilization: 0.1 distributed_executor_backend: "mp" - max_num_batched_tokens: 1000000 + max_num_batched_tokens: 100000 hf_config_name: thinker_config engine_input_source: [1] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index 8ecb4b0f8..c8637c8bd 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -34,6 +34,7 @@ def _make_runner(multimodal_outputs): runner.use_async_scheduling = False runner.device = torch.device("cpu") runner.supports_mm_inputs = False + runner.speculative_config = None return runner diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py index 83671b713..b68418e17 100644 --- a/vllm_omni/benchmarks/metrics/metrics.py +++ b/vllm_omni/benchmarks/metrics/metrics.py @@ -2,10 +2,10 @@ import warnings from dataclasses import dataclass import numpy as np -from transformers import PreTrainedTokenizerBase from vllm.benchmarks.datasets import SampleRequest from vllm.benchmarks.lib.endpoint_request_func import RequestFuncOutput from vllm.benchmarks.serve import MILLISECONDS_TO_SECONDS_CONVERSION, TERM_PLOTLIB_AVAILABLE, BenchmarkMetrics, TaskType +from vllm.tokenizers import TokenizerLike @dataclass @@ -134,7 +134,7 @@ def calculate_metrics( input_requests: list[SampleRequest], outputs: list[RequestFuncOutput], dur_s: float, - tokenizer: PreTrainedTokenizerBase, + tokenizer: TokenizerLike | None, selected_percentiles: list[float], goodput_config_dict: dict[str, float], task_type, @@ -169,17 +169,21 @@ def calculate_metrics( audio_rtfs: list[float] = [] audio_duration: list[float] = [] audio_frames: list[int] = [] + input_audio_duration = 0.0 for i in range(len(outputs)): if outputs[i].success: output_len = outputs[i].output_tokens if not output_len: - # We use the tokenizer to count the number of output tokens - # for some serving backends instead of looking at - # len(outputs[i].itl) since multiple output tokens may be - # bundled together - # Note : this may inflate the output token count slightly - output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids) + if tokenizer is None: + output_len = 1 + else: + # We use the tokenizer to count the number of output tokens + # for some serving backends instead of looking at + # len(outputs[i].itl) since multiple output tokens may be + # bundled together + # Note : this may inflate the output token count slightly + output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids) actual_output_lens.append(output_len) total_input += input_requests[i].prompt_len tpot = 0 @@ -196,6 +200,7 @@ def calculate_metrics( audio_duration.append(getattr(outputs[i], "audio_duration", 0.0)) audio_frames.append(getattr(outputs[i], "audio_frames", 0.0)) e2els.append(outputs[i].latency) + input_audio_duration += outputs[i].input_audio_duration completed += 1 else: actual_output_lens.append(0) @@ -335,6 +340,7 @@ def calculate_metrics( percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000) for p in selected_percentiles], max_output_tokens_per_s=max_output_tokens_per_s, max_concurrent_requests=max_concurrent_requests, + rtfx=input_audio_duration / dur_s, ) print_metrics( task_type, diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index c0a1fdba0..78792e81c 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -16,7 +16,6 @@ from typing import Literal import aiohttp from pydub import AudioSegment from tqdm.asyncio import tqdm -from transformers import PreTrainedTokenizerBase from vllm.benchmarks import datasets from vllm.benchmarks.datasets import SampleRequest from vllm.benchmarks.lib.endpoint_request_func import ( @@ -31,6 +30,7 @@ from vllm.benchmarks.lib.endpoint_request_func import ( _validate_api_url, ) from vllm.logger import init_logger +from vllm.tokenizers import TokenizerLike logger = init_logger(__name__) from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset @@ -134,6 +134,8 @@ async def async_request_openai_chat_omni_completions( messages = handler.add_chunk(chunk_bytes) for message in messages: + if type(message) is bytes: + message = message.decode("utf-8") # NOTE: SSE comments (often used as pings) start with # a colon. These are not JSON data payload and should # be skipped. @@ -311,7 +313,7 @@ async def benchmark( base_url: str, model_id: str, model_name: str, - tokenizer: PreTrainedTokenizerBase, + tokenizer: TokenizerLike | None, input_requests: list[SampleRequest], logprobs: int | None, request_rate: float, @@ -576,6 +578,7 @@ async def benchmark( "errors": [output.error for output in outputs], "max_output_tokens_per_s": metrics.max_output_tokens_per_s, "max_concurrent_requests": metrics.max_concurrent_requests, + "rtfx": metrics.rtfx, } else: result = { diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index a1d9e7765..fa5483b31 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -101,6 +101,7 @@ class OmniModelConfig(ModelConfig): def __post_init__( self, # Multimodal config init vars + language_model_only: bool, limit_mm_per_prompt: dict[str, int | dict[str, int]] | None, enable_mm_embeds: bool | None, media_io_kwargs: dict[str, dict[str, Any]] | None, @@ -117,6 +118,7 @@ class OmniModelConfig(ModelConfig): ) -> None: # Call parent's __post_init__ to handle all standard ModelConfig initialization super().__post_init__( + language_model_only=language_model_only, limit_mm_per_prompt=limit_mm_per_prompt, enable_mm_embeds=enable_mm_embeds, media_io_kwargs=media_io_kwargs, diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 82894236e..540fc391c 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -246,7 +246,7 @@ class OmniARScheduler(VLLMScheduler): generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else [] scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id) - if scheduled_spec_token_ids: + if scheduled_spec_token_ids and generated_token_ids: num_draft_tokens = len(scheduled_spec_token_ids) num_accepted = len(generated_token_ids) - 1 num_rejected = num_draft_tokens - num_accepted @@ -481,10 +481,7 @@ class OmniARScheduler(VLLMScheduler): assert request.is_finished() # 1. Standard cleanup parts from base _free_request - connector_delay_free_blocks = False - kv_xfer_params = None - if self.connector is not None: - connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) + connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) self.encoder_cache_manager.free(request) request_id = request.request_id diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index b85229a5d..a397e6085 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -6,6 +6,7 @@ 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.kv_cache_manager import KVCacheBlocks +from vllm.v1.core.sched.interface import PauseState from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.request_queue import create_request_queue from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler @@ -41,8 +42,12 @@ class OmniGenerationScheduler(VLLMScheduler): """ token_budget = self.max_num_scheduled_tokens + if self._pause_state == PauseState.PAUSED_ALL: + token_budget = 0 scheduled_timestamp = time.monotonic() + self.kv_cache_manager.new_step_starts() + scheduled_new_reqs: list[Request] = [] req_to_new_blocks: dict[str, KVCacheBlocks] = {} @@ -108,6 +113,8 @@ class OmniGenerationScheduler(VLLMScheduler): req_to_new_blocks[request.request_id] = new_blocks num_scheduled_tokens[request.request_id] = num_new_tokens cached_prompt_token_ids[request.request_id] = request.prompt_token_ids + if request.num_cached_tokens < 0: + request.num_cached_tokens = num_computed_tokens cached_additional_information[request.request_id] = getattr(request, "additional_information", None) token_budget -= num_new_tokens scheduled_running_reqs.append(request) @@ -119,7 +126,12 @@ class OmniGenerationScheduler(VLLMScheduler): # Fast path selection and scheduling (treat all as diffusion requests, # independent of pooling_params) - while self.waiting and token_budget > 0 and len(self.running) < self.max_num_running_reqs: + while ( + self.waiting + and token_budget > 0 + and len(self.running) < self.max_num_running_reqs + and self._pause_state == PauseState.UNPAUSED + ): request = self.waiting.peek_request() # OMNI: Skip requests that are not in self.requests if request.request_id not in self.requests or ( @@ -172,6 +184,8 @@ class OmniGenerationScheduler(VLLMScheduler): req_to_new_blocks[request.request_id] = new_blocks num_scheduled_tokens[request.request_id] = num_new_tokens + if request.num_cached_tokens < 0: + request.num_cached_tokens = 0 token_budget -= num_new_tokens scheduled_new_reqs.append(request) diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index f09e75c2b..c7a8c07c5 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -3,30 +3,26 @@ from dataclasses import dataclass, field from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput from vllm.v1.request import Request -from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload +from vllm_omni.engine import AdditionalInformationPayload @dataclass class OmniNewRequestData(NewRequestData): """New request data for omni models with embeddings support. - Extends NewRequestData to include prompt embeddings and additional - information for direct transfer between pipeline stages. + Extends NewRequestData to include additional information for direct + transfer between pipeline stages. + + Note: prompt_embeds is inherited from NewRequestData + (torch.Tensor | None). Args: - prompt_embeds: Optional serialized prompt embeddings payload - (overrides parent's torch.Tensor type with PromptEmbedsPayload - for cross-process serialization) external_req_id: Optional external request ID for tracking additional_information: Optional serialized additional information dictionary containing tensors or lists """ - # Optional serialized prompt embeddings (override parent type for serialization) - prompt_embeds: PromptEmbedsPayload | None = None # type: ignore[assignment] - # Optional external request ID for tracking external_req_id: str | None = None - # Optional serialized additional information additional_information: AdditionalInformationPayload | None = None @classmethod diff --git a/vllm_omni/diffusion/utils/hf_utils.py b/vllm_omni/diffusion/utils/hf_utils.py index 7935df1e4..6beb1823c 100644 --- a/vllm_omni/diffusion/utils/hf_utils.py +++ b/vllm_omni/diffusion/utils/hf_utils.py @@ -53,10 +53,8 @@ def is_diffusion_model(model_name: str) -> bool: # Strategy 2: Check using vllm's utility (works for both local and remote models) try: - # Try to get model_index.json using vllm's utility config_dict = get_hf_file_to_dict("model_index.json", model_name) - # Verify it has the required fields for a diffusers model - if config_dict.get("_class_name") and config_dict.get("_diffusers_version"): + if config_dict is not None and config_dict.get("_class_name") and config_dict.get("_diffusers_version"): logger.debug("Detected diffusion model via model_index.json") return True except Exception as e: diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index 1d1dd6f0f..91a0146c9 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -35,12 +35,26 @@ def try_send_via_connector( try: t0 = time.time() + # Strip non-serializable multimodal feature fields from original_prompt + # before including it in metadata. After stage-0 runs, the TokPrompt + # returned by render_chat_async may carry processed multimodal features + # (mm_kwargs, mm_placeholders, mm_hashes) that contain MultiModalKwargsItems + # objects, which are not supported by OmniMsgpackEncoder. The receiving + # side (try_recv_via_connector) only extracts "engine_inputs" from the + # payload and never uses "original_prompt", so stripping these fields + # only affects debug metadata and is safe. + _MM_FEATURE_KEYS = frozenset({"mm_kwargs", "mm_placeholders", "mm_hashes"}) + if isinstance(original_prompt, dict) and any(k in original_prompt for k in _MM_FEATURE_KEYS): + safe_prompt = {k: v for k, v in original_prompt.items() if k not in _MM_FEATURE_KEYS} + else: + safe_prompt = original_prompt + # Prepare data for connector payload_data = { "engine_inputs": next_inputs, "sampling_params": sampling_params, "metadata": { - "original_prompt": original_prompt, + "original_prompt": safe_prompt, "stage_transition": f"{stage_id}->{next_stage_id}", "timestamp": time.time(), }, diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 09984cf99..5c7384c1f 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -48,6 +48,7 @@ class SharedMemoryConnector(OmniConnectorBase): payload = self.serialize_obj(data) size = len(payload) + # Currently, we always use SHM. if True: # Use Shared Memory lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock" diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index c7289d8b8..db8d5dbaa 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -39,7 +39,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs self.connector = self.create_connector(model_config) super().__init__(model_config) - self.model_mode = getattr(model_config, "worker_type", "ar") + self.model_mode = getattr(model_config, "worker_type", None) or "ar" # State specific to Chunk management self.custom_process_next_stage_input_func = None custom_process_next_stage_input_func = getattr(model_config, "custom_process_next_stage_input_func", None) diff --git a/vllm_omni/distributed/omni_connectors/utils/serialization.py b/vllm_omni/distributed/omni_connectors/utils/serialization.py index 74aa6d4d0..ae9dc4953 100644 --- a/vllm_omni/distributed/omni_connectors/utils/serialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/serialization.py @@ -145,7 +145,7 @@ class OmniMsgpackEncoder: "encoder_prompt": obj.encoder_prompt, "encoder_prompt_token_ids": obj.encoder_prompt_token_ids, "num_cached_tokens": obj.num_cached_tokens, - "multi_modal_placeholders": obj.multi_modal_placeholders, + "multi_modal_placeholders": getattr(obj, "multi_modal_placeholders", None), "kv_transfer_params": obj.kv_transfer_params, } # Handle dynamically added multimodal_output attribute @@ -301,15 +301,22 @@ class OmniMsgpackDecoder: """Decode dict to RequestOutput. RequestOutput is not a dataclass, so msgspec.convert doesn't work. - We construct it manually, passing all known fields via **kwargs. + We construct it manually using only the known __init__ parameters to + avoid triggering the "Ignoring extra arguments" warning in vllm. + Fields that are not part of RequestOutput.__init__ (e.g. + multi_modal_placeholders, multimodal_output) are extracted first and + then restored as dynamic attributes after construction. """ - # Extract multimodal_output before constructing (it's dynamically added) + # Extract dynamically-added / non-init fields before constructing so + # they are not passed as unknown **kwargs to RequestOutput.__init__. mm_output = obj.pop("multimodal_output", None) + multi_modal_placeholders = obj.pop("multi_modal_placeholders", None) - # RequestOutput.__init__ accepts **kwargs for forward compatibility ro = RequestOutput(**obj) - # Restore dynamically added multimodal_output attribute + # Restore dynamic attributes that are not part of __init__. + if multi_modal_placeholders is not None: + setattr(ro, "multi_modal_placeholders", multi_modal_placeholders) if mm_output is not None: setattr(ro, "multimodal_output", mm_output) return ro diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index 0adfdb8d1..c4d42f95d 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -56,19 +56,19 @@ class AdditionalInformationPayload(msgspec.Struct): class OmniEngineCoreRequest(EngineCoreRequest): """Engine core request for omni models with embeddings support. - Extends the base EngineCoreRequest with support for prompt embeddings - and additional information payloads, enabling direct transfer of - pre-computed embeddings between pipeline stages. + Extends the base EngineCoreRequest with support for additional + information payloads, enabling direct transfer of pre-computed data + between pipeline stages. + + Note: prompt_embeds is inherited from EngineCoreRequest + (torch.Tensor | None). PromptEmbedsPayload should be decoded to + torch.Tensor before constructing this request. Attributes: - prompt_embeds: Optional serialized prompt embeddings payload for - direct transfer between stages additional_information: Optional serialized additional information dictionary containing tensors or lists to pass along with the request """ - # Optional prompt embeddings (direct-transfer version) - prompt_embeds: PromptEmbedsPayload | None = None # Optional additional information dictionary (serialized) additional_information: AdditionalInformationPayload | None = None diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index cbb42829f..c4f719703 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -116,6 +116,7 @@ class OmniEngineArgs(EngineArgs): # Keep compatibility when async args are constructed from partial payloads. limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {}) + language_model_only = getattr(self, "language_model_only", False) enable_mm_embeds = getattr(self, "enable_mm_embeds", False) interleave_mm_strings = getattr(self, "interleave_mm_strings", False) media_io_kwargs = getattr(self, "media_io_kwargs", {}) @@ -169,6 +170,7 @@ class OmniEngineArgs(EngineArgs): skip_tokenizer_init=self.skip_tokenizer_init, enable_prompt_embeds=self.enable_prompt_embeds, served_model_name=self.served_model_name, + language_model_only=language_model_only, limit_mm_per_prompt=limit_mm_per_prompt, enable_mm_embeds=enable_mm_embeds, interleave_mm_strings=interleave_mm_strings, @@ -276,6 +278,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): # Keep compatibility when async args are constructed from partial payloads. limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {}) + language_model_only = getattr(self, "language_model_only", False) enable_mm_embeds = getattr(self, "enable_mm_embeds", False) interleave_mm_strings = getattr(self, "interleave_mm_strings", False) media_io_kwargs = getattr(self, "media_io_kwargs", {}) @@ -329,6 +332,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): skip_tokenizer_init=self.skip_tokenizer_init, enable_prompt_embeds=self.enable_prompt_embeds, served_model_name=self.served_model_name, + language_model_only=language_model_only, limit_mm_per_prompt=limit_mm_per_prompt, enable_mm_embeds=enable_mm_embeds, interleave_mm_strings=interleave_mm_strings, diff --git a/vllm_omni/engine/input_processor.py b/vllm_omni/engine/input_processor.py index 99dc71cc4..5bbd16b38 100644 --- a/vllm_omni/engine/input_processor.py +++ b/vllm_omni/engine/input_processor.py @@ -1,6 +1,6 @@ import time from collections.abc import Mapping -from typing import Any, cast +from typing import Any import numpy as np import torch @@ -9,26 +9,14 @@ from vllm.inputs import ProcessorInputs, PromptType from vllm.inputs.parse import split_enc_dec_inputs from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry -from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalUUIDDict - -try: - from vllm.multimodal.processing import set_request_id -except ImportError: # vllm without set_request_id (older releases) - from contextlib import contextmanager - - @contextmanager - def set_request_id(_request_id: str): - yield - - +from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.multimodal.utils import argsort_mm_positions -from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams -from vllm.renderers.inputs import DictPrompt, TokPrompt +from vllm.renderers import BaseRenderer from vllm.sampling_params import SamplingParams from vllm.tasks import SupportedTask from vllm.utils import length_from_prompt_token_ids_or_embeds -from vllm.utils.torch_utils import set_default_torch_num_threads +from vllm.utils.jsontree import json_iter_leaves from vllm.v1.engine.input_processor import InputProcessor from vllm_omni.engine import ( @@ -42,6 +30,35 @@ from vllm_omni.lora.request import LoRARequest logger = init_logger(__name__) +_OMNI_EXTRA_KEYS = ( + "additional_information", + "prompt_embeds", + "negative_prompt", + "negative_prompt_embeds", +) + + +def reinject_omni_fields( + results: list[ProcessorInputs], + original_prompts: list[dict], +) -> None: + """Re-inject omni-specific fields that the upstream renderer discards. + + The upstream renderer's ``process_for_engine`` creates new dicts that only + copy standard vLLM fields (prompt_token_ids, multi_modal_data, …). + Omni-specific fields such as ``additional_information`` and + ``prompt_embeds`` are silently dropped. This helper copies them back from + the *original* parsed prompts into the renderer outputs so they survive + into ``OmniInputProcessor.process_inputs()``. + """ + for result, orig in zip(results, original_prompts): + if not isinstance(orig, dict): + continue + for key in _OMNI_EXTRA_KEYS: + val = orig.get(key) + if val is not None and key not in result: + result[key] = val + class OmniInputProcessor(InputProcessor): """Processor for omni models, handling multimodal inputs and embeddings. @@ -88,28 +105,29 @@ class OmniInputProcessor(InputProcessor): def __init__( self, vllm_config: VllmConfig, + renderer: BaseRenderer | None = None, + *, mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, ): - super().__init__(vllm_config, mm_registry) + super().__init__(vllm_config, renderer=renderer, mm_registry=mm_registry) self.input_preprocessor = OmniInputPreprocessor( - self.model_config, - vllm_config.observability_config, - mm_registry, - mm_processor_cache=self.mm_processor_cache, + vllm_config=vllm_config, + renderer=self.renderer, + mm_registry=mm_registry, ) def process_inputs( self, request_id: str, - prompt: PromptType | DictPrompt | TokPrompt, + prompt: PromptType | ProcessorInputs, params: SamplingParams | PoolingParams, + supported_tasks: tuple[SupportedTask, ...] = ("generate",), arrival_time: float | None = None, lora_request: LoRARequest | None = None, tokenization_kwargs: dict[str, Any] | None = None, trace_headers: Mapping[str, str] | None = None, priority: int = 0, data_parallel_rank: int | None = None, - supported_tasks: tuple[SupportedTask, ...] | None = None, resumable: bool = False, ) -> OmniEngineCoreRequest: """Process input prompt into an engine core request. @@ -123,6 +141,7 @@ class OmniInputProcessor(InputProcessor): request_id: Unique identifier for this request prompt: Input prompt (text, token IDs, embeddings, or multimodal) params: Sampling or pooling parameters for generation + supported_tasks: Tuple of supported tasks for validation arrival_time: Optional arrival timestamp (defaults to current time) lora_request: Optional LoRA adapter request tokenization_kwargs: Optional additional tokenization arguments @@ -130,19 +149,17 @@ class OmniInputProcessor(InputProcessor): priority: Request priority (higher values processed first) data_parallel_rank: Optional data parallel rank for distributed inference + resumable: Whether the request supports streaming input Returns: - Tuple of (prompt_string, OmniEngineCoreRequest) where: - - prompt_string: The original prompt as a string, or None if - using embeddings - - OmniEngineCoreRequest: Processed request ready for the engine + OmniEngineCoreRequest ready for the engine Raises: ValueError: If data_parallel_rank is out of range or prompt_embeds has incorrect shape """ - self._validate_lora(lora_request) self._validate_params(params, supported_tasks) + self._validate_lora(lora_request) parallel_config = self.vllm_config.parallel_config dp_size = parallel_config.data_parallel_size @@ -151,50 +168,23 @@ class OmniInputProcessor(InputProcessor): if data_parallel_rank is not None and not (0 <= data_parallel_rank < num_ranks): raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {num_ranks}).") - if arrival_time is None: - arrival_time = time.time() - - # Optionally generate multimodal hash overrides to avoid hashing - # multimodal data items by their content as their identifiers. - - # NOTE: when users explicitly turn off BOTH prefix caching and input - # processing caching, no multimodal features or embeddings will be - # reused across requests, therefore identifying multimodal data items - # by their content is no longer necessary, and we create uuids with - # request id-modality-index as multimodal hash overrides. - if ( - self.model_config.multimodal_config - and self.model_config.multimodal_config.mm_processor_cache_gb == 0 - and not self.cache_config.enable_prefix_caching - ): - mm_uuids = self._maybe_build_mm_uuids(request_id, prompt) + # Short-circuit for prompts already processed by the renderer + # (they carry a "type" key). Raw prompts must still go through the + # omni preprocessor which preserves additional_information, etc. + if isinstance(prompt, dict) and "type" in prompt: + if arrival_time is None: + arrival_time = prompt.get("arrival_time", time.time()) + processed_inputs: ProcessorInputs = prompt # type: ignore[assignment] else: - # Otherwise, use user-provided uuids as multimodal hash overrides - # if provided. - self._validate_mm_uuids(prompt) - if isinstance(prompt, dict): - mm_uuids = cast(MultiModalUUIDDict | None, prompt.get("multi_modal_uuids")) - else: - mm_uuids = None + if arrival_time is None: + arrival_time = time.time() - # Process inputs, which includes: - # 1. Tokenize text prompt, with LoRA request if one exists. - # 2. For multimodal models with a merged preprocessor, preprocess - # multimodal data and expand prompt token ids accordingly. - with set_request_id(request_id), set_default_torch_num_threads(): - processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess( + processed_inputs = self.input_preprocessor.preprocess( prompt, tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, ) - current_platform.validate_request( - prompt=prompt, - params=params, - processed_inputs=processed_inputs, - ) - - eos_token_id = self.input_preprocessor.get_eos_token_id() + self._platform_validate_request(processed_inputs, params) encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs) self._validate_model_inputs(encoder_inputs, decoder_inputs) @@ -216,7 +206,10 @@ class OmniInputProcessor(InputProcessor): if sampling_params.max_tokens is None: seq_len = length_from_prompt_token_ids_or_embeds(prompt_token_ids, prompt_embeds) sampling_params.max_tokens = self.model_config.max_model_len - seq_len - sampling_params.update_from_generation_config(self.generation_config_fields, eos_token_id) + sampling_params.update_from_generation_config( + self.generation_config_fields, + self.renderer.get_eos_token_id(), + ) if self.tokenizer is not None: sampling_params.update_from_tokenizer(self.tokenizer) else: @@ -230,6 +223,13 @@ class OmniInputProcessor(InputProcessor): decoder_mm_positions = decoder_inputs["mm_placeholders"] decoder_mm_hashes = decoder_inputs["mm_hashes"] + if not all(isinstance(leaf, str) for leaf in json_iter_leaves(decoder_mm_hashes)): + raise ValueError( + f"mm_hashes must contain only strings, got: {decoder_mm_hashes}. " + "This is likely due to an incorrect custom implementation of " + "MultiModalProcessor.apply method." + ) + # Merge and flatten multimodal placeholders, hashes and inputs # from dictionaries to lists, and sort them by each item's position # in the input sequence. @@ -281,7 +281,6 @@ class OmniInputProcessor(InputProcessor): mm_features=mm_features, sampling_params=sampling_params, pooling_params=pooling_params, - eos_token_id=eos_token_id, arrival_time=arrival_time, lora_request=lora_request, cache_salt=decoder_inputs.get("cache_salt"), diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index a37e84097..43d02e85b 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -199,7 +199,7 @@ class OmniRequestState(RequestState): if not ( finished or self.sent_tokens_offset == 0 - or len(self.detokenizer.output_token_ids) - self.sent_tokens_offset >= self.stream_interval + or self.detokenizer.num_output_tokens() - self.sent_tokens_offset >= self.stream_interval ): return None @@ -207,7 +207,7 @@ class OmniRequestState(RequestState): # Send tokens from the offset in DELTA mode, otherwise all # tokens are sent. new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :] - self.sent_tokens_offset = len(self.detokenizer.output_token_ids) + self.sent_tokens_offset = self.detokenizer.num_output_tokens() external_req_id = self.external_req_id output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts) @@ -264,8 +264,10 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): def __init__( self, tokenizer: TokenizerLike | None, + *, log_stats: bool, stream_interval: int = 1, + tracing_enabled: bool = False, engine_core_output_type: str | None = None, ): """Initialize the multimodal output processor. @@ -278,7 +280,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): (e.g., "image", "audio", "latent"). Used to tag multimodal outputs with the correct modality key. """ - super().__init__(tokenizer=tokenizer, log_stats=log_stats, stream_interval=stream_interval) + super().__init__( + tokenizer=tokenizer, + log_stats=log_stats, + stream_interval=stream_interval, + tracing_enabled=tracing_enabled, + ) self.engine_core_output_type = engine_core_output_type def add_request( @@ -320,8 +327,6 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): log_stats=self.log_stats, stream_interval=self.stream_interval, ) - if self._requests_drained.is_set(): - self._requests_drained.clear() self.request_states[request_id] = req_state if parent_req: self.parent_requests[parent_req.request_id] = parent_req diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index f7e95aa32..46bbd0c19 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -432,7 +432,10 @@ class AsyncOmni(OmniBase): ) -> AsyncGenerator[OmniRequestOutput, None]: all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)} submit_flag = True + _loop_iter = 0 + _last_progress_ts = time.time() while not all(all_stages_finished.values()): + _loop_iter += 1 for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]): if all_stages_finished[stage_id]: continue @@ -441,6 +444,7 @@ class AsyncOmni(OmniBase): except asyncio.QueueEmpty: await asyncio.sleep(0.001) continue + _last_progress_ts = time.time() engine_outputs, finished, output_to_yield = self._process_single_result( result, stage, @@ -449,11 +453,20 @@ class AsyncOmni(OmniBase): ) if submit_flag and stage_id == 0: submit_flag = False - prompt_token_ids = engine_outputs.prompt_token_ids + prompt_token_ids = getattr(engine_outputs, "prompt_token_ids", None) + if prompt_token_ids is None: + prompt_token_ids = [] engine_input = copy.deepcopy(prompt) - next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) + try: + next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) + except Exception: + raise engine_input["prompt_token_ids"] = [0] * next_prompt_len engine_input["multi_modal_data"] = engine_input["mm_processor_kwargs"] = None + for _mm_key in ("mm_kwargs", "mm_hashes", "mm_placeholders", "multi_modal_uuids"): + engine_input.pop(_mm_key, None) + if engine_input.get("type") == "multimodal": + engine_input["type"] = "token" for i in range(1, len(self.stage_list)): task = { "request_id": request_id, @@ -706,6 +719,16 @@ class AsyncOmni(OmniBase): stage.submit(abort_task) return None + async def get_supported_tasks(self) -> set[str]: + tasks: set[str] = set() + has_comprehension = any(stage.is_comprehension for stage in self.stage_list) + if has_comprehension: + tasks.add("generate") + for stage in self.stage_list: + if getattr(stage, "final_output_type", None) == "audio": + tasks.add("speech") + return tasks if tasks else {"generate"} + async def get_vllm_config(self) -> VllmConfig: for stage in self.stage_list: if stage.is_comprehension: @@ -744,6 +767,8 @@ class AsyncOmni(OmniBase): OMNI: Required by upstream OpenAIServingModels.__init__ which accesses engine_client.renderer. """ + if self.input_processor is None: + return None return self.input_processor.renderer async def do_log_stats(self) -> None: diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py index 02d175ad0..d63e4f154 100644 --- a/vllm_omni/entrypoints/async_omni_llm.py +++ b/vllm_omni/entrypoints/async_omni_llm.py @@ -9,7 +9,7 @@ import torch from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry -from vllm.tokenizers import cached_tokenizer_from_config +from vllm.renderers import renderer_from_config from vllm.tracing import init_tracer from vllm.transformers_utils.config import maybe_register_config_serialize_by_value from vllm.usage.usage_lib import UsageContext @@ -106,21 +106,17 @@ class AsyncOmniLLM(AsyncLLM): "enabling logging without default stat loggers" ) - if self.model_config.skip_tokenizer_init: - tokenizer = None - else: - # Tokenizer (+ ensure liveness if running in another process). - tokenizer = cached_tokenizer_from_config(model_config=vllm_config.model_config) - # InputProcessor (converts Inputs --> EngineCoreRequests). self.input_processor = OmniInputProcessor( vllm_config=vllm_config, mm_registry=mm_registry, ) + self.renderer = renderer_from_config(self.vllm_config) + # OutputProcessor (converts EngineCoreOutputs --> RequestOutput). self.output_processor = MultimodalOutputProcessor( - tokenizer=tokenizer, + tokenizer=self.renderer.tokenizer, log_stats=self.log_stats, engine_core_output_type=engine_args.engine_output_type, ) diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index ec9248e30..27f470b4b 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -1,4 +1,4 @@ -from collections.abc import Callable +from collections.abc import Callable, Sequence from typing import Any import cloudpickle @@ -8,9 +8,11 @@ from tqdm import tqdm # External library imports (vLLM) from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field from vllm.entrypoints.llm import LLM +from vllm.inputs import ProcessorInputs, PromptType from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput, RequestOutput from vllm.plugins.io_processors import get_io_processor +from vllm.renderers.inputs.preprocess import parse_model_prompt from vllm.usage.usage_lib import UsageContext from vllm.utils.counter import Counter from vllm.v1.engine.llm_engine import LLMEngine @@ -19,7 +21,7 @@ from vllm_omni.distributed.omni_connectors import initialize_orchestrator_connec # Internal imports (our code) from vllm_omni.engine.arg_utils import OmniEngineArgs -from vllm_omni.engine.input_processor import OmniInputProcessor +from vllm_omni.engine.input_processor import OmniInputProcessor, reinject_omni_fields from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.entrypoints.utils import ( filter_dataclass_kwargs, @@ -174,12 +176,50 @@ class OmniLLM(LLM): self.supported_tasks = supported_tasks + # Keep parity with vLLM's LLM initialization fields so inherited + # generate/chat preprocessing paths work as expected. + self.renderer = self.llm_engine.renderer # Load the Input/Output processor plugin if any io_processor_plugin = self.llm_engine.model_config.io_processor_plugin self.io_processor = get_io_processor(self.llm_engine.vllm_config, io_processor_plugin) self.model_config = self.llm_engine.model_config self.input_processor = self.llm_engine.input_processor + # Parity with upstream LLM for pooling/classify entrypoints + chat_template = kwargs.get("chat_template", None) + from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template + from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors + + self.chat_template = load_chat_template(chat_template) + self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template) + self.init_pooling_io_processors = init_pooling_io_processors( + supported_tasks=supported_tasks, + model_config=self.model_config, + renderer=self.renderer, + chat_template_config=self.chat_template_config, + ) + + # ------------------------------------------------------------------ + # Override upstream _preprocess_cmpl so that omni-specific fields + # (additional_information, prompt_embeds, …) survive the renderer's + # process_for_engine step which only copies standard vLLM keys. + # ------------------------------------------------------------------ + + def _preprocess_cmpl( + self, + prompts: Sequence[PromptType], + tokenization_kwargs: dict[str, Any] | None = None, + ) -> Sequence[ProcessorInputs]: + renderer = self.renderer + model_config = self.model_config + + parsed_prompts = [parse_model_prompt(model_config, prompt) for prompt in prompts] + tok_params = renderer.default_cmpl_tok_params.with_kwargs(**(tokenization_kwargs or {})) + results = renderer.render_cmpl(parsed_prompts, tok_params) + + reinject_omni_fields(results, parsed_prompts) + return results + def close(self) -> None: """Close resources. @@ -197,7 +237,9 @@ class OmniLLM(LLM): except Exception as e: logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True) - def _run_engine(self, *, use_tqdm: bool | Callable[..., tqdm] = True) -> list[RequestOutput | PoolingRequestOutput]: + def _run_engine( + self, output_type=None, *, use_tqdm: bool | Callable[..., tqdm] = True + ) -> list[RequestOutput | PoolingRequestOutput]: # Initialize tqdm. if use_tqdm: num_requests = self.llm_engine.get_num_unfinished_requests() diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 1ce9dae48..52e54f22f 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -76,20 +76,16 @@ def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int = """ from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available - nvml_available = is_process_scoped_memory_available() - pid_host = detect_pid_host() - - if nvml_available and pid_host: - logger.info( + if is_process_scoped_memory_available() and detect_pid_host(): + logger.debug( "NVML process-scoped memory available and PID host is available — concurrent init is safe, skipping locks" ) yield return else: - logger.info( - "Using sequential init locks (nvml_available=%s, pid_host=%s)", - nvml_available, - pid_host, + logger.debug( + "NVML unavailable or PID host is not available (usually inside a container, " + "--pid=host is not set in docker run command) — using sequential init locks" ) from vllm_omni.platforms import current_omni_platform @@ -1356,7 +1352,7 @@ async def _stage_worker_async( usage_context=usage_context, engine_args=omni_engine_args, disable_log_stats=bool( - engine_args.get("disable_log_stats", True) or getattr(omni_engine_args, "disable_log_stats", True) + engine_args.get("disable_log_stats", False) or getattr(omni_engine_args, "disable_log_stats", False) ), ) if hasattr(stage_engine, "log_stats") and stage_engine.log_stats: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 2134411df..6db97dd2d 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -287,6 +287,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, ssl_certfile=args.ssl_certfile, ssl_ca_certs=args.ssl_ca_certs, ssl_cert_reqs=args.ssl_cert_reqs, + ssl_ciphers=args.ssl_ciphers, h11_max_incomplete_event_size=args.h11_max_incomplete_event_size, h11_max_header_count=args.h11_max_header_count, **uvicorn_kwargs, diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index e2021ebc5..88a5378e2 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -453,6 +453,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): default_template_content_format, ).with_defaults(default_template_kwargs) + # OMNI: When use_audio_in_video=True, the qwen2_5_omni_thinker mm + # processor asserts that audio items are present alongside video items + # (inside render_chat_async). We must inject audio_url items into the + # messages BEFORE calling render_chat_async so the mm processor can + # count them correctly during tokenisation. + mm_proc_kw = getattr(request, "mm_processor_kwargs", None) or {} + if mm_proc_kw.get("use_audio_in_video", False): + messages = await self._inject_audio_from_video_urls(messages) + (conversation,), (engine_prompt,) = await renderer.render_chat_async( [messages], chat_params, @@ -462,28 +471,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): }, ) - # OMNI: When use_audio_in_video=True, the upstream renderer does not - # extract audio from video. We do it here after rendering so that the - # audio data is present in multi_modal_data before the engine processes - # the request. - mm_proc_kw = getattr(request, "mm_processor_kwargs", None) or {} - if mm_proc_kw.get("use_audio_in_video", False) and isinstance(engine_prompt, dict): - mm_data = engine_prompt.get("multi_modal_data") - if mm_data is not None and "video" in mm_data and "audio" not in mm_data: - from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async - - video_urls: list[str] = [] - for msg in messages: - for part in msg.get("content") or []: - if isinstance(part, dict) and part.get("type") == "video_url": - url = part.get("video_url", {}).get("url") - if url: - video_urls.append(url) - - if video_urls: - audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls)) - engine_prompt.setdefault("multi_modal_data", {})["audio"] = list(audios) - tokenizer = renderer.get_tokenizer() # tool parsing is done only if a tool_parser has been set and if @@ -533,6 +520,95 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): return conversation, [engine_prompt] + async def _inject_audio_from_video_urls( + self, + messages: list[ChatCompletionMessageParam], + ) -> list[ChatCompletionMessageParam]: + """Pre-extract audio from video URLs and inject as audio_url content items. + + When use_audio_in_video=True, the qwen2_5_omni_thinker multimodal + processor requires that the number of audio items equals the number of + video items (it subtracts mm_counts["video"] from mm_counts["audio"]). + The client only sends video_url items; this method adds the matching + audio_url items on the server side before the renderer processes them. + """ + import io + + from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async + + new_messages: list[ChatCompletionMessageParam] = [] + for msg in messages: + content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None) + if not isinstance(content, list): + new_messages.append(msg) + continue + + video_urls = [ + part.get("video_url", {}).get("url") + for part in content + if isinstance(part, dict) and part.get("type") == "video_url" and part.get("video_url", {}).get("url") + ] + + if not video_urls: + new_messages.append(msg) + continue + + audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls)) + + audio_items: list[dict] = [] + for audio_array, sample_rate in audios: + buf = io.BytesIO() + if soundfile is not None: + soundfile.write(buf, audio_array, samplerate=int(sample_rate), format="WAV") + else: + import struct + + import numpy as np + + audio_np = np.asarray(audio_array, dtype=np.float32) + sr = int(sample_rate) + num_channels = 1 + bits_per_sample = 32 + num_frames = len(audio_np) + data_size = num_frames * num_channels * (bits_per_sample // 8) + # Write minimal RIFF/WAV header + buf.write(b"RIFF") + buf.write(struct.pack(" list[SamplingParams]: final_sampling_params_list = [] for sampling_params in sampling_params_list: @@ -1213,11 +1289,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): latest_delta_len = len(delta_message.tool_calls[0].function.arguments) # get the expected call based on partial JSON - # parsing which "autocompletes" the JSON - expected_call = json.dumps( - tool_parser.prev_tool_call_arr[index].get("arguments", {}), - ensure_ascii=False, - ) + # parsing which "autocompletes" the JSON. + # Tool parsers (e.g. Qwen3Coder) store + # arguments as a JSON string in + # prev_tool_call_arr. Calling json.dumps() + # on an already-serialized string would + # double-serialize it (e.g. '{"k":1}' becomes + # '"{\\"k\\":1}"'), which then causes the + # replace() below to fail and append the + # entire double-serialized string as a + # spurious final delta. + args = tool_parser.prev_tool_call_arr[index].get("arguments", {}) + if isinstance(args, str): + expected_call = args + else: + expected_call = json.dumps(args, ensure_ascii=False) # get what we've streamed so far for arguments # for the current tool diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 429bb0f5e..15f31627f 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -4,7 +4,7 @@ from typing_extensions import assert_never from vllm.inputs.data import EmbedsInputs, SingletonInputs from vllm.inputs.preprocess import InputPreprocessor from vllm.logger import init_logger -from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict +from vllm.multimodal.inputs import MultiModalInputs from vllm.renderers.inputs import SingletonDictPrompt from vllm_omni.inputs.data import ( @@ -30,8 +30,6 @@ class OmniInputPreprocessor(InputPreprocessor): self, parsed_content: OmniTextPrompt, tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: MultiModalUUIDDict | None = None, ) -> OmniTokenInputs | MultiModalInputs: """Process text prompts with support for mm_processor_kwargs. @@ -49,9 +47,7 @@ class OmniInputPreprocessor(InputPreprocessor): multi_modal_data, mm_processor_kwargs, tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, ) - # Preserve prompt_embeds and additional_information prompt_embeds = parsed_content.get("prompt_embeds") if prompt_embeds is not None: inputs["prompt_embeds"] = prompt_embeds @@ -59,14 +55,11 @@ class OmniInputPreprocessor(InputPreprocessor): if additional_information is not None: inputs["additional_information"] = additional_information elif mm_processor_kwargs: - # Support mm_processor_kwargs without multi_modal_data - # Use case: GLM-Image text-to-image needs processor to generate grid tokens inputs = self._process_multimodal( prompt_text, - {}, # Empty multi_modal_data + {}, mm_processor_kwargs, tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, ) else: prompt_token_ids = self._tokenize_prompt( @@ -78,7 +71,12 @@ class OmniInputPreprocessor(InputPreprocessor): prompt_embeds=parsed_content.get("prompt_embeds"), additional_information=parsed_content.get("additional_information"), ) - + prompt_embeds = parsed_content.get("prompt_embeds") + if prompt_embeds is not None: + inputs["prompt_embeds"] = prompt_embeds + additional_information = parsed_content.get("additional_information") + if additional_information is not None: + inputs["additional_information"] = additional_information if cache_salt := parsed_content.get("cache_salt"): inputs["cache_salt"] = cache_salt @@ -88,8 +86,6 @@ class OmniInputPreprocessor(InputPreprocessor): self, parsed_content: OmniTokensPrompt, tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: MultiModalUUIDDict | None = None, ) -> OmniTokenInputs | MultiModalInputs: prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs) prompt_embeds = parsed_content.get("prompt_embeds") @@ -102,21 +98,23 @@ class OmniInputPreprocessor(InputPreprocessor): inputs = self._process_multimodal( prompt_token_ids, multi_modal_data, - parsed_content.get("mm_processor_kwargs") or {}, + parsed_content.get("mm_processor_kwargs"), tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, + mm_uuids=parsed_content.get("multi_modal_uuids"), ) - if prompt_embeds is not None: - inputs["prompt_embeds"] = prompt_embeds - if additional_information is not None: - inputs["additional_information"] = additional_information + else: inputs = token_inputs_omni( prompt_token_ids=prompt_token_ids, prompt_embeds=prompt_embeds, additional_information=additional_information, ) - + if prompt_embeds is not None: + inputs["prompt_embeds"] = prompt_embeds + if additional_information is not None: + inputs["additional_information"] = additional_information + if prompt_text := parsed_content.get("prompt"): + inputs["prompt"] = prompt_text if cache_salt := parsed_content.get("cache_salt"): inputs["cache_salt"] = cache_salt @@ -145,8 +143,6 @@ class OmniInputPreprocessor(InputPreprocessor): self, prompt: SingletonDictPrompt, tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: MultiModalUUIDDict | None = None, ) -> SingletonInputs: """ Extract the singleton inputs from a prompt. @@ -159,20 +155,18 @@ class OmniInputPreprocessor(InputPreprocessor): * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance """ + if "prompt_embeds" in prompt: + return self._process_embeds(prompt) # type: ignore[arg-type] + if "prompt_token_ids" in prompt: return self._process_tokens( prompt, # type: ignore[arg-type] - mm_uuids=mm_uuids, ) - if "prompt_embeds" in prompt: - return self._process_embeds(prompt) # type: ignore[arg-type] - if "prompt" in prompt: return self._process_text( prompt, # type: ignore[arg-type] tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, ) assert_never(prompt) # type: ignore[arg-type] diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index ac7d8120d..3362ae17f 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -19,7 +19,6 @@ from vllm.model_executor.models.interfaces import ( from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniThinkerDummyInputsBuilder, ) -from vllm.model_executor.models.qwen3_omni_moe_thinker import Qwen3Omni_VisionTransformer from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, @@ -33,6 +32,7 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mt ) from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( Qwen3MoeLLMForCausalLM, + Qwen3Omni_VisionTransformer, Qwen3OmniMoeConditionalGenerationMixin, Qwen3OmniMoeThinkerMultiModalProcessor, Qwen3OmniMoeThinkerProcessingInfo, diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index 8c67852e9..7dfb988eb 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -43,9 +43,12 @@ from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import ( from transformers.models.whisper import WhisperFeatureExtractor from vllm.compilation.decorators import support_torch_compile from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.distributed import get_pp_group +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size from vllm.inputs.data import PromptType from vllm.logger import init_logger +from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, +) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.interfaces import ( @@ -70,7 +73,9 @@ from vllm.model_executor.models.qwen2_audio import Qwen2AudioProcessingInfo from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM from vllm.model_executor.models.qwen3_moe import Qwen3MoeModel as _Qwen3MoeLLMModel from vllm.model_executor.models.qwen3_omni_moe_thinker import ( - Qwen3Omni_VisionTransformer, + Qwen3Omni_VisionTransformer as _Qwen3Omni_VisionTransformer, +) +from vllm.model_executor.models.qwen3_omni_moe_thinker import ( Qwen3OmniMoeAudioEncoder, _get_feat_extract_output_lengths, ) @@ -104,6 +109,90 @@ except (ImportError, ModuleNotFoundError): logger = init_logger(__name__) + +class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer): + """Subclass that fixes Qwen2_5_VisionAttention.forward() compatibility. + + The upstream Qwen3_VisionBlock.forward() does not pass the + ``sequence_lengths`` argument required by the updated + Qwen2_5_VisionAttention.forward() signature. This subclass overrides + ``forward()`` to compute ``sequence_lengths`` via MMEncoderAttention + (following the pattern in qwen3_vl.py) and calls block internals + directly so the argument is forwarded correctly. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.tp_size = get_tensor_model_parallel_world_size() + + def forward( + self, + x: torch.Tensor, + grid_thw, + ) -> torch.Tensor: + hidden_states = x.to(device=self.device, dtype=self.dtype) + hidden_states = self.patch_embed(hidden_states) + + if self.apply_vit_abs_pos_embed: + pos_embeds = self.fast_pos_embed_interpolate(grid_thw) + hidden_states = hidden_states + pos_embeds + rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw) + + if isinstance(grid_thw, torch.Tensor): + grid_thw_np = grid_thw.cpu().numpy().astype(np.int32) + else: + grid_thw_np = np.array(grid_thw, dtype=np.int32) + + cu_seqlens = np.repeat(grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]).cumsum(axis=0, dtype=np.int32) + cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens]) + + sequence_lengths = MMEncoderAttention.maybe_compute_sequence_lengths(self.attn_backend, cu_seqlens) + if sequence_lengths is not None: + sequence_lengths = torch.from_numpy(sequence_lengths).to(self.device, non_blocking=True) + max_seqlen = torch.tensor( + MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens), + dtype=torch.int32, + device=self.device, + ) + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens( + self.attn_backend, + cu_seqlens, + self.hidden_size, + self.tp_size, + ) + cu_seqlens = torch.from_numpy(cu_seqlens).to(self.device, non_blocking=True) + + hidden_states = hidden_states.unsqueeze(1) + + hidden_states_list = [] + deepstack_visual_indexes = self.deepstack_visual_indexes + + for layer_num, blk in enumerate(self.blocks): + hidden_states = hidden_states + blk.attn( + blk.norm1(hidden_states), + cu_seqlens=cu_seqlens, + rotary_pos_emb_cos=rotary_pos_emb_cos, + rotary_pos_emb_sin=rotary_pos_emb_sin, + max_seqlen=max_seqlen, + sequence_lengths=sequence_lengths, + ) + hidden_states = hidden_states + blk.mlp(blk.norm2(hidden_states)) + + if deepstack_visual_indexes is not None and layer_num in deepstack_visual_indexes: + hidden_states_list.append(hidden_states) + + hidden_states = self.merger(hidden_states) + + if deepstack_visual_indexes is not None: + processed_hidden_states_list = [hidden_states] + for idx, x_ds in enumerate(hidden_states_list): + x_ds = self.merger_list[idx](x_ds) + processed_hidden_states_list.append(x_ds) + hidden_states = torch.cat(processed_hidden_states_list, dim=1) + + return hidden_states + + # Speech input languages supported by Qwen3-Omni # From: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct ISO639_1_SUPPORTED_LANGS = { @@ -228,6 +317,38 @@ class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProc def get_supported_mm_limits(self) -> Mapping[str, int | None]: return {"audio": None, "image": None, "video": None} + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int] | None = None, + ) -> Mapping[str, int] | None: + mm_counts = mm_counts or {} + requested_modalities = {m for m, c in mm_counts.items() if c > 0} + mm_max_tokens: dict[str, int] = {} + + if requested_modalities & {"image", "video"}: + vl_tokens = Qwen2_5_VLProcessingInfo.get_mm_max_tokens_per_item( + self, + seq_len=seq_len, + mm_counts=mm_counts, + ) + mm_max_tokens.update({m: vl_tokens[m] for m in ["image", "video"] if m in requested_modalities}) + + if "audio" in requested_modalities: + audio_tokens = Qwen2AudioProcessingInfo.get_mm_max_tokens_per_item( + self, + seq_len=seq_len, + mm_counts=mm_counts, + ) + if audio_tokens is None: + feature_extractor = self.get_feature_extractor() + max_audio_samples = feature_extractor.chunk_length * feature_extractor.sampling_rate + max_audio_tokens = int(_get_feat_extract_output_lengths(torch.tensor([max_audio_samples])).item()) + audio_tokens = {"audio": max_audio_tokens} + mm_max_tokens["audio"] = audio_tokens["audio"] + + return mm_max_tokens + Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder @@ -253,7 +374,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( return x # NOTE: WhisperFeatureExtractor cannot handle empty list of audios - feature_extractor = self.info.get_feature_extractor() + feature_extractor = self.info.get_feature_extractor(**mm_kwargs) hop_length = feature_extractor.hop_length if audios: # NOTE: Qwen3-Omni processor accept "audio" @@ -653,8 +774,7 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix feature_lens=audio_feature_lengths, aftercnn_lens=audio_output_lengths, ) - # OMNI: audio_tower.forward() returns hidden_states tensor directly - audio_features = audio_outputs + audio_features = audio_outputs if isinstance(audio_outputs, torch.Tensor) else audio_outputs.last_hidden_state return audio_features.split(audio_output_lengths.tolist()) diff --git a/vllm_omni/request.py b/vllm_omni/request.py index a869fa23c..3ec325316 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -82,7 +82,6 @@ class OmniRequest(Request): mm_features=request.mm_features, sampling_params=request.sampling_params, pooling_params=request.pooling_params, - eos_token_id=request.eos_token_id, arrival_time=request.arrival_time, lora_request=request.lora_request, cache_salt=request.cache_salt, diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index cf261429f..1ea56f724 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -22,7 +22,9 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( ) from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output +from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( @@ -258,6 +260,9 @@ class GPUARModelRunner(OmniGPUModelRunner): # Run the model. # Use persistent buffers for CUDA graphs. + # When spec decode is enabled, delay clearing connector metadata + # until after draft model runs in sample_tokens. + clear_kv_metadata = self.speculative_config is None with ( set_forward_context( attn_metadata, @@ -270,7 +275,9 @@ class GPUARModelRunner(OmniGPUModelRunner): slot_mapping=slot_mappings, # OMNI: required for KV cache operations ), record_function_or_nullcontext("gpu_model_runner: forward"), - self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, + self.maybe_get_kv_connector_output( + scheduler_output, clear_metadata=clear_kv_metadata + ) as kv_connector_output, ): model_output = self._model_forward( input_ids=input_ids, @@ -376,13 +383,12 @@ class GPUARModelRunner(OmniGPUModelRunner): self, grammar_output: GrammarOutput | None, ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: - kv_connector_output = self.kv_connector_output - self.kv_connector_output = None - kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None) self.kv_extracted_req_ids = None if self.execute_model_state is None: + kv_connector_output = self.kv_connector_output + self.kv_connector_output = None # Nothing to do (PP non-final rank case), output isn't used. if not kv_connector_output: return None # type: ignore[return-value] @@ -454,10 +460,14 @@ class GPUARModelRunner(OmniGPUModelRunner): spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens <= self.effective_drafter_max_model_len ) - if spec_config.use_eagle() and not spec_config.disable_padded_drafter_batch: - # EAGLE speculative decoding can use the GPU sampled tokens - # as inputs, and does not need to wait for bookkeeping to finish. - assert isinstance(self.drafter, EagleProposer) + use_gpu_toks = ( + spec_config.use_eagle() or spec_config.uses_draft_model() or spec_config.uses_extract_hidden_states() + ) and not spec_config.disable_padded_drafter_batch + if use_gpu_toks: + assert isinstance( + self.drafter, + EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer, + ) sampled_token_ids = sampler_output.sampled_token_ids if input_fits_in_drafter: propose_draft_token_ids(sampled_token_ids) @@ -503,9 +513,19 @@ class GPUARModelRunner(OmniGPUModelRunner): # tokens on the CPU, so they are run after bookkeeping. propose_draft_token_ids(valid_sampled_token_ids) + # Clear KV connector metadata after draft model runs (if spec decode). + # This was deferred from target model forward to allow draft model + # to also save its KV cache. + if self.speculative_config is not None: + self.clear_kv_connector_metadata() + with record_function_or_nullcontext("gpu_model_runner: eplb"): self.eplb_step() + # kv_connector_output may be modified during drafting + kv_connector_output = self.kv_connector_output + self.kv_connector_output = None + hidden_states_cpu = hidden_states.detach().to("cpu").contiguous() num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None) if num_scheduled_tokens_np is None: diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 67e530190..05785d7aa 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -23,7 +23,9 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output +from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( EMPTY_MODEL_RUNNER_OUTPUT, @@ -263,6 +265,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # Run the model. # Use persistent buffers for CUDA graphs. + # When spec decode is enabled, delay clearing connector metadata + # until after draft model runs in sample_tokens. + clear_kv_metadata = self.speculative_config is None with ( set_forward_context( attn_metadata, @@ -275,7 +280,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): slot_mapping=slot_mappings, # OMNI: required for KV cache operations ), record_function_or_nullcontext("Forward"), - self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, + self.maybe_get_kv_connector_output( + scheduler_output, clear_metadata=clear_kv_metadata + ) as kv_connector_output, ): outputs = self._run_generation_model( input_ids=input_ids, @@ -344,6 +351,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): ) = self.execute_model_state self.execute_model_state = None + # Clear KV connector metadata after draft model runs (if spec decode). + if self.speculative_config is not None: + self.clear_kv_connector_metadata() + pooler_output: list[object] = [] if isinstance(multimodal_outputs, torch.Tensor): assert multimodal_outputs.shape[0] == 1, ( @@ -487,7 +498,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): create_mixed_batch: If True, create a mixed batch with both decode (1 token) and prefill (multiple tokens) requests. remove_lora: If False, dummy LoRAs are not destroyed after the run - num_active_loras: Number of active LoRAs to capture for. + num_active_loras: Number of distinct active LoRAs to capture for. + LoRA is activated when num_active_loras > 0. """ mm_config = self.vllm_config.model_config.multimodal_config if mm_config and mm_config.mm_encoder_only: @@ -515,7 +527,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # Set num_scheduled_tokens based on num_tokens and max_num_seqs # for dummy run with LoRA so that the num_reqs collectively # has num_tokens in total. - assert num_tokens <= self.scheduler_config.max_num_batched_tokens + assert num_tokens <= self.max_num_tokens max_num_reqs = self.scheduler_config.max_num_seqs if create_mixed_batch: assert not uniform_decode @@ -566,7 +578,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # activated later in the context manager, but we need to know the # LoRA state when determining the batch descriptor for capture force_has_lora=num_active_loras > 0, - # Capture shape specialization for specific active LoRA counts. + # `force_num_active_loras` is used for cudagraph capture; because we + # need to capture graphs for specific num_active_loras counts force_num_active_loras=num_active_loras, ) ) @@ -604,33 +617,36 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): ubatch_slices=ubatch_slices_padded, ) - # If force_attention is True, we always capture attention. Otherwise, - # it only happens for cudagraph_runtime_mode=FULL. - if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: - if create_mixed_batch: - # In the mixed batch mode (used for FI warmup), we use - # shorter sequence lengths to run faster. - # TODO(luka) better system for describing dummy batches - seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] - else: - seq_lens = max_query_len # type: ignore[assignment] - self.seq_lens.np[:num_reqs] = seq_lens - self.seq_lens.np[num_reqs:] = 0 - self.seq_lens.copy_to_gpu() + with self.synchronize_input_prep(): + # If force_attention is True, we always capture attention. + # Otherwise, it only happens for cudagraph_runtime_mode=FULL. + if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: + if create_mixed_batch: + # In the mixed batch mode (used for FI warmup), we use + # shorter sequence lengths to run faster. + # TODO(luka) better system for describing dummy batches + seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] + else: + seq_lens = max_query_len # type: ignore[assignment] + self.seq_lens.np[:num_reqs] = seq_lens + self.seq_lens.np[num_reqs:] = 0 + self.seq_lens.copy_to_gpu() - cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) - self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens - self.query_start_loc.copy_to_gpu() + cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() - pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL - attn_metadata, _ = self._build_attention_metadata( - num_tokens=num_tokens_unpadded, - num_reqs=num_reqs_padded, - max_query_len=max_query_len, - ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, - for_cudagraph_capture=is_graph_capturing, - slot_mappings=slot_mappings_by_group, - ) + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded if pad_attn else None, + num_reqs=num_reqs_padded, + max_query_len=max_query_len, + ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices), + for_cudagraph_capture=is_graph_capturing, + slot_mappings=slot_mappings_by_group, + use_spec_decode=self.speculative_config is not None, + ) with self.maybe_dummy_run_with_lora( self.lora_config, @@ -720,8 +736,16 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): else: hidden_states = outputs hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) - if self.speculative_config and self.speculative_config.use_eagle(): - assert isinstance(self.drafter, EagleProposer) + if self.speculative_config and ( + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + or self.speculative_config.uses_extract_hidden_states() + ): + assert isinstance( + self.drafter, + EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer, + ) + assert self.speculative_config is not None # Eagle currently only supports PIECEWISE cudagraphs. # Therefore only use cudagraphs if the main model uses PIECEWISE # NOTE(lucas): this is a hack, need to clean up. diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 966d6e5b8..dc353d8d4 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -13,7 +13,9 @@ from vllm.sampling_params import SamplingType from vllm.tracing import instrument from vllm.utils.import_utils import LazyLoader from vllm.utils.math_utils import cdiv +from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices @@ -275,6 +277,11 @@ class OmniGPUModelRunner(GPUModelRunner): # Add new requests to the cached states. for new_req_data in scheduler_output.scheduled_new_reqs: req_id = new_req_data.req_id + if req_id in self.requests: + req_state = self._update_streaming_request(req_id, new_req_data) + reqs_to_add.append(req_state) + continue + sampling_params = new_req_data.sampling_params pooling_params = new_req_data.pooling_params @@ -316,10 +323,7 @@ class OmniGPUModelRunner(GPUModelRunner): arr = np.frombuffer(payload.data, dtype=dtype) arr = arr.reshape(payload.shape) pe_cpu = torch.from_numpy(arr) - # Store temporarily on CPU; later moved to device in builder setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) - # Also replace payload with Tensor for user visibility in - # scheduler_output try: new_req_data.prompt_embeds = pe_cpu # type: ignore[assignment] except Exception: @@ -350,7 +354,6 @@ class OmniGPUModelRunner(GPUModelRunner): info_dict[k] = entry.list_data if info_dict: self.model_intermediate_buffer[req_id] = info_dict - # Backward compatible: mirror to old setattr location setattr( self.requests[req_id], "additional_information_cpu", @@ -358,7 +361,6 @@ class OmniGPUModelRunner(GPUModelRunner): ) except Exception as e: logger.error(f"Error decoding additional information: {e}") - pass if sampling_params and sampling_params.prompt_logprobs is not None: self.num_prompt_logprobs[req_id] = ( @@ -421,18 +423,15 @@ class OmniGPUModelRunner(GPUModelRunner): req_state.num_computed_tokens = num_computed_tokens if not is_last_rank: - # When using PP, the scheduler sends the sampled tokens back, - # because there's no direct communication between the first- - # stage worker and the last-stage worker. - new_token_ids = req_data.new_token_ids[i] - # Add the sampled token(s) from the previous step (if any). - # This doesn't include "unverified" tokens like spec tokens. - num_new_tokens = num_computed_tokens + len(new_token_ids) - req_state.num_tokens - if num_new_tokens == 1: - # Avoid slicing list in most common case. - req_state.output_token_ids.append(new_token_ids[-1]) - elif num_new_tokens > 0: - req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:]) + if not req_data.new_token_ids: + new_token_ids: list[int] = [] + else: + new_token_ids = req_data.new_token_ids[i] + num_new_tokens = num_computed_tokens + len(new_token_ids) - req_state.num_tokens + if num_new_tokens == 1: + req_state.output_token_ids.append(new_token_ids[-1]) + elif num_new_tokens > 0: + req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:]) elif num_output_tokens < len(req_state.output_token_ids): # Some output tokens were discarded due to a sync-KV-load # failure. Align the cached state. @@ -533,7 +532,6 @@ class OmniGPUModelRunner(GPUModelRunner): remove_lora: bool = True, is_graph_capturing: bool = False, num_active_loras: int = 0, - activate_lora: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Run a dummy forward pass to warm up/profile run or capture the @@ -556,12 +554,9 @@ class OmniGPUModelRunner(GPUModelRunner): create_mixed_batch: If True, create a mixed batch with both decode (1 token) and prefill (multiple tokens) requests. remove_lora: If False, dummy LoRAs are not destroyed after the run - num_active_loras: Number of active LoRAs to capture for. - activate_lora: Backward-compatible override for LoRA activation. + num_active_loras: Number of distinct active LoRAs to capture for. + LoRA is activated when num_active_loras > 0. """ - if activate_lora is None: - activate_lora = num_active_loras > 0 - mm_config = self.vllm_config.model_config.multimodal_config if mm_config and mm_config.mm_encoder_only: # The current dummy run only covers LM execution, so we can skip it. @@ -588,7 +583,7 @@ class OmniGPUModelRunner(GPUModelRunner): # Set num_scheduled_tokens based on num_tokens and max_num_seqs # for dummy run with LoRA so that the num_reqs collectively # has num_tokens in total. - assert num_tokens <= self.scheduler_config.max_num_batched_tokens + assert num_tokens <= self.max_num_tokens max_num_reqs = self.scheduler_config.max_num_seqs if create_mixed_batch: assert not uniform_decode @@ -638,8 +633,9 @@ class OmniGPUModelRunner(GPUModelRunner): # `force_has_lora` is used for cudagraph capture; because LoRA is # activated later in the context manager, but we need to know the # LoRA state when determining the batch descriptor for capture - force_has_lora=activate_lora, - # Capture shape specialization for specific active LoRA counts. + force_has_lora=num_active_loras > 0, + # `force_num_active_loras` is used for cudagraph capture; because we + # need to capture graphs for specific num_active_loras counts force_num_active_loras=num_active_loras, ) ) @@ -676,40 +672,43 @@ class OmniGPUModelRunner(GPUModelRunner): ubatch_slices=ubatch_slices_padded, ) - # If force_attention is True, we always capture attention. Otherwise, - # it only happens for cudagraph_runtime_mode=FULL. - if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: - if create_mixed_batch: - # In the mixed batch mode (used for FI warmup), we use - # shorter sequence lengths to run faster. - # TODO(luka) better system for describing dummy batches - seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] - else: - seq_lens = max_query_len # type: ignore[assignment] - self.seq_lens.np[:num_reqs] = seq_lens - self.seq_lens.np[num_reqs:] = 0 - self.seq_lens.copy_to_gpu() + with self.synchronize_input_prep(): + # If force_attention is True, we always capture attention. + # Otherwise, it only happens for cudagraph_runtime_mode=FULL. + if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: + if create_mixed_batch: + # In the mixed batch mode (used for FI warmup), we use + # shorter sequence lengths to run faster. + # TODO(luka) better system for describing dummy batches + seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] + else: + seq_lens = max_query_len # type: ignore[assignment] + self.seq_lens.np[:num_reqs] = seq_lens + self.seq_lens.np[num_reqs:] = 0 + self.seq_lens.copy_to_gpu() - cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) - self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens - self.query_start_loc.copy_to_gpu() + cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() - pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL - attn_metadata, _ = self._build_attention_metadata( - num_tokens=num_tokens_unpadded, - num_reqs=num_reqs_padded, - max_query_len=max_query_len, - ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, - for_cudagraph_capture=is_graph_capturing, - slot_mappings=slot_mappings_by_group, - ) + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded if pad_attn else None, + num_reqs=num_reqs_padded, + max_query_len=max_query_len, + ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices), + for_cudagraph_capture=is_graph_capturing, + slot_mappings=slot_mappings_by_group, + use_spec_decode=self.speculative_config is not None, + ) with self.maybe_dummy_run_with_lora( self.lora_config, num_scheduled_tokens, num_sampled_tokens, - activate_lora, remove_lora, + num_active_loras, ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens @@ -798,8 +797,16 @@ class OmniGPUModelRunner(GPUModelRunner): else: hidden_states = outputs hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) - if self.speculative_config and self.speculative_config.use_eagle(): - assert isinstance(self.drafter, EagleProposer) + if self.speculative_config and ( + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + or self.speculative_config.uses_extract_hidden_states() + ): + assert isinstance( + self.drafter, + EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer, + ) + assert self.speculative_config is not None # Eagle currently only supports PIECEWISE cudagraphs. # Therefore only use cudagraphs if the main model uses PIECEWISE # NOTE(lucas): this is a hack, need to clean up. @@ -812,7 +819,7 @@ class OmniGPUModelRunner(GPUModelRunner): # lora cases when cudagraph_specialize_lora is enabled. This is a # short term mitigation for issue mentioned in # https://github.com/vllm-project/vllm/issues/28334 - if self.compilation_config.cudagraph_specialize_lora and activate_lora: + if self.compilation_config.cudagraph_specialize_lora and num_active_loras > 0: use_cudagraphs = False self.drafter.dummy_run( @@ -847,64 +854,100 @@ class OmniGPUModelRunner(GPUModelRunner): logit_indices_device = torch.from_numpy(logit_indices).to(self.device, non_blocking=True) return hidden_states, hidden_states[logit_indices_device] - def _decode_and_store_request_payloads(self, scheduler_output: "SchedulerOutput") -> None: - """Decode per-request prompt_embeds and additional_information for newly - scheduled requests and store them to CPU in the request state. - This version avoids hard dependency on payload classes by duck-typing.""" + # ------------------------------------------------------------------ + # Payload decoding helpers (torch.Tensor passthrough + legacy + # PromptEmbedsPayload / AdditionalInformationPayload support) + # ------------------------------------------------------------------ + + @staticmethod + def _resolve_prompt_embeds_cpu( + pe: "torch.Tensor | object | None", + ) -> torch.Tensor | None: + """Convert *prompt_embeds* to a contiguous CPU tensor. + + Accepts: + - ``torch.Tensor`` – moved to CPU as-is (the normal path after + upstream added ``prompt_embeds`` to ``EngineCoreRequest``). + - Legacy ``PromptEmbedsPayload`` (or any duck-typed object with + ``.data``, ``.shape``, ``.dtype``) – decoded via numpy. + - ``None`` – returns ``None``. + """ + if pe is None: + return None try: - new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) - if not new_reqs: - return - for nr in new_reqs: - req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None) - if req_id is None: - continue - # prompt_embeds - payload_pe = getattr(nr, "prompt_embeds", None) - pe_cpu = None - if payload_pe is not None: - if isinstance(payload_pe, torch.Tensor): - pe_cpu = payload_pe.detach().to("cpu").contiguous() - else: - # Try duck-typing a payload with data/shape/dtype - data = getattr(payload_pe, "data", None) - shape = getattr(payload_pe, "shape", None) - if data is not None and shape is not None: - dt = np.dtype(getattr(payload_pe, "dtype", "float32")) - arr = np.frombuffer(data, dtype=dt) - arr = arr.reshape(shape) - pe_cpu = torch.from_numpy(arr.copy()) - if pe_cpu is not None and req_id in self.requests: - setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) - # additional_information - payload_info = getattr(nr, "additional_information", None) - if payload_info is not None: - logger.warning_once( - "additional_information on request data is deprecated, use model_intermediate_buffer" - ) - info_dict = {} - if isinstance(payload_info, dict): - info_dict = payload_info - else: - # Try duck-typing a payload with entries, each entry may have - # tensor_data/tensor_dtype/tensor_shape or list_data - entries = getattr(payload_info, "entries", None) - if isinstance(entries, dict): - for k, entry in entries.items(): - tensor_data = getattr(entry, "tensor_data", None) - if tensor_data is not None: - dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) - arr = np.frombuffer(tensor_data, dtype=dt) - arr = arr.reshape(getattr(entry, "tensor_shape", ())) - info_dict[k] = torch.from_numpy(arr.copy()) - else: - info_dict[k] = getattr(entry, "list_data", None) - if info_dict and req_id in self.requests: - self.model_intermediate_buffer[req_id] = info_dict - # Backward compatible: mirror to old setattr location - setattr(self.requests[req_id], "additional_information_cpu", info_dict) - except Exception as e: - logger.error(f"Error decoding prompt_embeds / additional_information: {e}") + if isinstance(pe, torch.Tensor): + return pe.detach().cpu().contiguous() + data = getattr(pe, "data", None) + shape = getattr(pe, "shape", None) + if data is not None and shape is not None: + dt = np.dtype(getattr(pe, "dtype", "float32")) + arr = np.frombuffer(data, dtype=dt).reshape(shape) + return torch.from_numpy(arr.copy()) + except Exception: + logger.exception("Failed to decode prompt_embeds payload") + return None + + @staticmethod + def _resolve_additional_information( + payload: "dict | object | None", + ) -> dict[str, object]: + """Convert an *additional_information* payload to a plain dict. + + Accepts: + - ``dict`` – returned as-is. + - ``AdditionalInformationPayload`` (or duck-typed with + ``.entries``) – decoded entry-by-entry. + - ``None`` – returns ``{}``. + """ + if payload is None: + return {} + if isinstance(payload, dict): + return payload + try: + entries = getattr(payload, "entries", None) + if not isinstance(entries, dict): + return {} + info: dict[str, object] = {} + for k, entry in entries.items(): + tensor_data = getattr(entry, "tensor_data", None) + if tensor_data is not None: + dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) + arr = np.frombuffer(tensor_data, dtype=dt) + arr = arr.reshape(getattr(entry, "tensor_shape", ())) + info[k] = torch.from_numpy(arr.copy()) + else: + info[k] = getattr(entry, "list_data", None) + return info + except Exception: + logger.exception("Failed to decode additional_information payload") + return {} + + def _decode_and_store_request_payloads( + self, + scheduler_output: "SchedulerOutput", + ) -> None: + """Decode per-request prompt_embeds and additional_information for + newly scheduled requests and store them on CPU in the request state. + """ + new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) + if not new_reqs: + return + for nr in new_reqs: + req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None) + if req_id is None or req_id not in self.requests: + continue + pe_cpu = self._resolve_prompt_embeds_cpu(getattr(nr, "prompt_embeds", None)) + if pe_cpu is not None: + setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) + info_payload = getattr(nr, "additional_information", None) + if info_payload is not None: + logger.warning_once( + "additional_information on request data is deprecated, use model_intermediate_buffer" + ) + info_dict = self._resolve_additional_information(info_payload) + if info_dict: + self.model_intermediate_buffer[req_id] = info_dict + setattr(self.requests[req_id], "additional_information_cpu", info_dict) def _gather_runtime_additional_information(self) -> list[dict]: """Gather per-request model_intermediate_buffer in batch order."""