diff --git a/docs/api/README.md b/docs/api/README.md index 258c3e7b5..3a110f491 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -53,8 +53,10 @@ Engine classes for offline and online inference. Core scheduling and caching components. +- [vllm_omni.core.sched.omni_ar_scheduler.KVCacheTransferData][] - [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][] - [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][] +- [vllm_omni.core.sched.output.OmniCachedRequestData][] - [vllm_omni.core.sched.output.OmniNewRequestData][] - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][] - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][] diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index cd9198228..d26246533 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -15,6 +15,12 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` +If you want to open async chunking for qwen3-omni, launch the server with command below + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml +``` + If you have custom stage configs file, launch the server with command below ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index 9c1f411f8..9db1060b7 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -12,6 +12,12 @@ Please refer to [README.md](../../../README.md) vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` +If you want to open async chunking for qwen3-omni, launch the server with command below + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml +``` + If you have custom stage configs file, launch the server with command below ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file diff --git a/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py index e4768a8aa..304b27c4e 100644 --- a/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py +++ b/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py @@ -1,4 +1,5 @@ import base64 +import concurrent.futures import os from typing import NamedTuple @@ -414,51 +415,67 @@ def run_multimodal_generation(args) -> None: else: output_modalities = None - chat_completion = client.chat.completions.create( - messages=[ - get_system_prompt(), - prompt, - ], - model=model_name, - modalities=output_modalities, - extra_body=extra_body, - stream=args.stream, - ) + # Test multiple concurrent completions + num_concurrent_requests = args.num_concurrent_requests + with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor: + # Submit multiple completion requests concurrently + futures = [ + executor.submit( + client.chat.completions.create, + messages=[ + get_system_prompt(), + prompt, + ], + model=model_name, + modalities=output_modalities, + extra_body=extra_body, + stream=args.stream, + ) + for _ in range(num_concurrent_requests) + ] + + # Wait for all requests to complete and collect results + chat_completions = [future.result() for future in concurrent.futures.as_completed(futures)] + + assert len(chat_completions) == num_concurrent_requests count = 0 if not args.stream: - for choice in chat_completion.choices: - if choice.message.audio: - audio_data = base64.b64decode(choice.message.audio.data) - audio_file_path = f"audio_{count}.wav" - with open(audio_file_path, "wb") as f: - f.write(audio_data) - print(f"Audio saved to {audio_file_path}") - count += 1 - elif choice.message.content: - print("Chat completion output from text:", choice.message.content) - else: - printed_content = False - for chunk in chat_completion: - for choice in chunk.choices: - if hasattr(choice, "delta"): - content = getattr(choice.delta, "content", None) - else: - content = None - - if getattr(chunk, "modality", None) == "audio" and content: - audio_data = base64.b64decode(content) + # Verify all completions succeeded + for chat_completion in chat_completions: + for choice in chat_completion.choices: + if choice.message.audio: + audio_data = base64.b64decode(choice.message.audio.data) audio_file_path = f"audio_{count}.wav" with open(audio_file_path, "wb") as f: f.write(audio_data) - print(f"\nAudio saved to {audio_file_path}") + print(f"Audio saved to {audio_file_path}") count += 1 + elif choice.message.content: + print("Chat completion output from text:", choice.message.content) + else: + printed_content = False + for chat_completion in chat_completions: + for chunk in chat_completion: + for choice in chunk.choices: + if hasattr(choice, "delta"): + content = getattr(choice.delta, "content", None) + else: + content = None - elif getattr(chunk, "modality", None) == "text": - if not printed_content: - printed_content = True - print("\ncontent:", end="", flush=True) - print(content, end="", flush=True) + if getattr(chunk, "modality", None) == "audio" and content: + audio_data = base64.b64decode(content) + audio_file_path = f"audio_{count}.wav" + with open(audio_file_path, "wb") as f: + f.write(audio_data) + print(f"\nAudio saved to {audio_file_path}") + count += 1 + + elif getattr(chunk, "modality", None) == "text": + if not printed_content: + printed_content = True + print("\ncontent:", end="", flush=True) + print(content, end="", flush=True) def parse_args(): @@ -517,6 +534,12 @@ def parse_args(): action="store_true", help="Stream the response.", ) + parser.add_argument( + "--num-concurrent-requests", + type=int, + default=1, + help="Number of concurrent requests to send. Default is 1.", + ) return parser.parse_args() diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index da0f5716b..4fe136288 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -1,4 +1,5 @@ import warnings +from dataclasses import field from typing import Any import torch @@ -42,6 +43,7 @@ class OmniModelConfig(ModelConfig): Attributes: stage_id: Identifier for the stage in a multi-stage pipeline (default: 0) + async_chunk: If set to True, perform async chunk model_stage: Stage type identifier, e.g., "thinker" or "talker" (default: "thinker") model_arch: Model architecture name @@ -49,6 +51,8 @@ class OmniModelConfig(ModelConfig): engine_output_type: Optional output type specification for the engine. Used to route outputs to appropriate processors (e.g., "image", "audio", "latents"). If None, output type is inferred. + stage_connector_config: Stage connector configuration dictionary. + Contains "name" (connector name), "extra" (extra connector config). Example: >>> config = OmniModelConfig( @@ -59,10 +63,18 @@ class OmniModelConfig(ModelConfig): """ stage_id: int = 0 + async_chunk: bool = False model_stage: str = "thinker" model_arch: str = "Qwen2_5OmniForConditionalGeneration" engine_output_type: str | None = None hf_config_name: str | None = None + custom_process_next_stage_input_func: str | None = None + stage_connector_config: dict[str, Any] = field( + default_factory=lambda: { + "name": "SharedMemoryConnector", + "extra": {}, + } + ) omni_kv_config: dict | None = None @property diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 46717cf04..03eb5ff21 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -1,5 +1,6 @@ from __future__ import annotations +import importlib from collections import defaultdict from dataclasses import asdict, dataclass from time import time @@ -19,6 +20,10 @@ from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm_omni.distributed.omni_connectors.adapter import get_chunk, put_chunk +from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory +from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec + logger = init_logger(__name__) @@ -59,6 +64,25 @@ class OmniARScheduler(VLLMScheduler): # Track requests that have already triggered prefill transfer to avoid duplicates self.transfer_triggered_requests: set[str] = set() + model_config = self.vllm_config.model_config + self.omni_connector = None + if model_config.async_chunk: + connector_config = model_config.stage_connector_config + connector_specs = ConnectorSpec( + name=connector_config.get("name", "SharedMemoryConnector"), + extra=connector_config.get("extra", {}), + ) + self.omni_connector = OmniConnectorFactory.create_connector(connector_specs) + + custom_process_next_stage_input_func = getattr( + self.vllm_config.model_config, "custom_process_next_stage_input_func", None + ) + if custom_process_next_stage_input_func: + module_path, func_name = custom_process_next_stage_input_func.rsplit(".", 1) + module = importlib.import_module(module_path) + self.custom_process_next_stage_input_func = getattr(module, func_name) + + self.stage_id = getattr(self.vllm_config.model_config, "stage_id", None) def _get_kv_transfer_criteria(self) -> dict | None: # Note: vllm_config is available in Scheduler after super().__init__ @@ -156,6 +180,8 @@ class OmniARScheduler(VLLMScheduler): new_list.append(omni_nr) scheduler_output.scheduled_new_reqs = new_list # type: ignore[assignment] + if self.omni_connector is not None: + get_chunk(self.omni_connector, scheduler_output) # Add information about requests needing KV cache transfer scheduler_output.finished_requests_needing_kv_transfer = self.get_finished_requests_needing_kv_transfer() @@ -337,6 +363,9 @@ class OmniARScheduler(VLLMScheduler): num_nans_in_logits=request.num_nans_in_logits, ) ) + if self.omni_connector is not None: + custom_process_next_stage_input_func = self.custom_process_next_stage_input_func + put_chunk(self.omni_connector, pooler_output, request, custom_process_next_stage_input_func) else: # Invariant: EngineCore returns no partial prefill outputs. assert not prompt_logprobs_tensors diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index d9a4b12d3..0164029d2 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -12,11 +12,27 @@ from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutp from vllm.v1.request import Request, RequestStatus from vllm.v1.spec_decode.metrics import SpecDecodingStats -from vllm_omni.core.sched.output import OmniNewRequestData +from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData +from vllm_omni.distributed.omni_connectors.adapter import get_chunk_for_generation +from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory +from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec from vllm_omni.outputs import OmniModelRunnerOutput class OmniGenerationScheduler(VLLMScheduler): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + model_config = self.vllm_config.model_config + self.omni_connector = None + if model_config.async_chunk: + connector_config = model_config.stage_connector_config + connector_specs = ConnectorSpec( + name=connector_config.get("name", "SharedMemoryConnector"), + extra=connector_config.get("extra", {}), + ) + self.omni_connector = OmniConnectorFactory.create_connector(connector_specs) + self.stage_id = getattr(self.vllm_config.model_config, "stage_id", None) + def schedule(self) -> SchedulerOutput: """Diffusion fast path: - Feed all input tokens of the request at once @@ -32,27 +48,54 @@ class OmniGenerationScheduler(VLLMScheduler): req_to_new_blocks: dict[str, KVCacheBlocks] = {} num_scheduled_tokens: dict[str, int] = {} + scheduled_running_reqs: list[Request] = [] scheduled_spec_decode_tokens: dict[str, list[int]] = {} scheduled_encoder_inputs: dict[str, list[int]] = {} + cached_prompt_token_ids: dict[str, list[int]] = {} # Temporary queue: preserve waiting order, do not disturb non-diffusion requests skipped_waiting_requests = create_request_queue(self.policy) + req_index = 0 + while req_index < len(self.running) and token_budget > 0: + request = self.running[req_index] + if self.omni_connector is not None: + get_chunk_for_generation(self.omni_connector, request) + num_computed_tokens = request.num_computed_tokens + required_tokens = max(len(request.prompt_token_ids) - num_computed_tokens, 1) + num_new_tokens = min(required_tokens, token_budget) + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, + num_lookahead_tokens=self.num_lookahead_tokens, + ) + if new_blocks is None: + # Allocation failed (e.g., VRAM pressure); stop fast path and + # fall back to default scheduling + # Put the current request back to the head of the waiting queue + # Note: the original queue order is preserved + break + if self.log_stats: + request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp) + 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 + token_budget -= num_new_tokens + scheduled_running_reqs.append(request) + req_index += 1 # 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: request = self.waiting.peek_request() + if self.omni_connector is not None: + get_chunk_for_generation(self.omni_connector, request) # Uniformly treat as diffusion. A feature flag can be added later # via config or request tag. # Allocate all input tokens for the request in one shot # (allocate 1 placeholder if zero) - required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1) - if required_tokens > token_budget: - # Insufficient budget to process all inputs at once; - # stop fast path attempt - break - num_new_tokens = required_tokens + required_tokens = max(len(request.prompt_token_ids), 1) + num_new_tokens = min(required_tokens, token_budget) new_blocks = self.kv_cache_manager.allocate_slots( request, num_new_tokens, @@ -68,7 +111,6 @@ class OmniGenerationScheduler(VLLMScheduler): # Officially schedule this request request = self.waiting.pop_request() self.running.append(request) - request.status = RequestStatus.RUNNING if self.log_stats: request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp) @@ -109,13 +151,24 @@ class OmniGenerationScheduler(VLLMScheduler): ] # No running/resumed reqs scheduled in our fast path cached_reqs_data = self._make_cached_request_data( - running_reqs=[], + running_reqs=scheduled_running_reqs, resumed_reqs=[], num_scheduled_tokens=num_scheduled_tokens, spec_decode_tokens=scheduled_spec_decode_tokens, req_to_new_blocks=req_to_new_blocks, ) + cached_reqs_data = OmniCachedRequestData( + req_ids=cached_reqs_data.req_ids, + resumed_req_ids=cached_reqs_data.resumed_req_ids, + new_token_ids=cached_reqs_data.new_token_ids, + all_token_ids=cached_reqs_data.all_token_ids, + new_block_ids=cached_reqs_data.new_block_ids, + num_computed_tokens=cached_reqs_data.num_computed_tokens, + num_output_tokens=cached_reqs_data.num_output_tokens, + prompt_token_ids=cached_prompt_token_ids, + ) + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, @@ -257,21 +310,20 @@ class OmniGenerationScheduler(VLLMScheduler): new_logprobs = None new_token_ids = generated_token_ids kv_transfer_params = None - status_before_stop = request.status pooler_output = None if pooler_outputs: pooler_output = pooler_outputs[req_index] # Diffusion request: completes in one step; mark finished and free resources - request.status = RequestStatus.FINISHED_STOPPED - # Optional: set a stop_reason for front-end clarity - # (does not affect protocol) - request.stop_reason = request.stop_reason # or "generation_done" - kv_transfer_params = self._free_request(request) - if status_before_stop == RequestStatus.RUNNING: + if request.status == RequestStatus.FINISHED_STOPPED or ( + self.omni_connector is None and request.num_computed_tokens >= request.num_prompt_tokens + ): + request.status = RequestStatus.FINISHED_STOPPED + # Optional: set a stop_reason for front-end clarity + # (does not affect protocol) + request.stop_reason = request.stop_reason # or "generation_done" + kv_transfer_params = self._free_request(request) stopped_running_reqs.add(request) - else: - stopped_preempted_reqs.add(request) # Extract sample logprobs if needed. if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs: diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index 205513ac9..3176ec14a 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -1,6 +1,6 @@ from dataclasses import dataclass -from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.core.sched.output import CachedRequestData, NewRequestData from vllm.v1.request import Request from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload @@ -53,3 +53,14 @@ class OmniNewRequestData(NewRequestData): prefill_token_ids=prefill_token_ids, additional_information=request.additional_information, ) + + +@dataclass +class OmniCachedRequestData(CachedRequestData): + """Cached request data for omni models with embeddings support. + + Args: + prompt_token_ids: Mapping from request ID to list of prompt token IDs + """ + + prompt_token_ids: dict[int, list[int]] diff --git a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py b/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py index 5fb0cc670..54327be81 100644 --- a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py @@ -207,10 +207,11 @@ class GPUDiffusionModelRunner: start_time = time.time() while True: + get_key = f"omni_{from_stage}_to_{to_stage}_kv_cache_{req.request_id}" result = self.connector.get( from_stage=from_stage, to_stage=to_stage, - request_id=f"kv_cache_{req.request_id}", + get_key=get_key, ) if result: break diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index cfc959017..da3abf4a9 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -5,7 +5,14 @@ import time from collections.abc import Callable -from typing import Any +from typing import TYPE_CHECKING, Any + +import torch +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.request import Request, RequestStatus + +if TYPE_CHECKING: + from .connectors.base import OmniConnectorBase from vllm_omni.entrypoints.stage_utils import OmniStageTaskType @@ -170,3 +177,225 @@ def try_recv_via_connector( # but for Stage-0 seed it should be there. # We'll return None to let caller handle error if strictly required. return None, None + + +def get_chunk( + connector: "OmniConnectorBase", + scheduler_output: SchedulerOutput, +) -> None: + """Retrieve a chunk of pooling output. + + Args: + connector: OmniConnectorBase instance + scheduler_output: Partial scheduler output dictionary + + Returns: + None: This function modifies scheduler_output in place + """ + stage_id = connector.stage_id + if stage_id == 0: + return + + target_stage_id = stage_id - 1 + # Handle new requests + for new_req_data in scheduler_output.scheduled_new_reqs: + req_id = new_req_data.req_id[0:25] + chunk_id = connector.get_requests[req_id] + connector_get_key = f"{req_id}_{target_stage_id}_{chunk_id}" + payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key) + if payload_data: + new_req_data.additional_information = payload_data + if payload_data.get("finished"): + connector.finished_requests.add(req_id) + + # Handle cached/running requests + cached_reqs = scheduler_output.scheduled_cached_reqs + if not hasattr(cached_reqs, "additional_information"): + cached_reqs.additional_information = {} + + for i, cached_req_id in enumerate(cached_reqs.req_ids): + req_id = cached_req_id[0:25] + if req_id in connector.finished_requests: + continue + chunk_id = connector.get_requests[req_id] + connector_get_key = f"{req_id}_{target_stage_id}_{chunk_id}" + payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key) + if payload_data: + cached_reqs.additional_information[cached_req_id] = payload_data + if payload_data.get("finished"): + connector.finished_requests.add(req_id) + + +def get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key): + # Wait for data from previous stage + import time + + # TODO: add correct check mechanism for the payload_data + max_wait = 300 + for _ in range(max_wait): + result = connector.get( + from_stage=str(target_stage_id), + to_stage=str(stage_id), + get_key=connector_get_key, + ) + payload_data = None + if result: + payload_data, size = result + if payload_data: + connector.request_prompt_token_ids[req_id] = payload_data.get("thinker_input_ids", []) + connector.get_requests[req_id] += 1 + logger.debug(f"[Stage-{stage_id}] Received one chunk for request {connector_get_key}") + break + time.sleep(0.1) + return payload_data + + +def get_chunk_for_generation( + connector: "OmniConnectorBase", + request: Request, +) -> None: + """Retrieve a chunk of pooling output. + + Args: + connector: OmniConnectorBase instance + request: Request object + + Returns: + None: This function modifies request in place + """ + stage_id = connector.stage_id + target_stage_id = stage_id - 1 + request_id = request.request_id[0:25] + + if request_id in connector.finished_requests: + return + + chunk_id = connector.get_requests[request_id] + connector_get_key = f"{request_id}_{target_stage_id}_{chunk_id}" + payload_data = get_through_connector(connector, target_stage_id, stage_id, request_id, connector_get_key) + if not payload_data: + return + + if payload_data.get("finished"): + connector.finished_requests.add(request_id) + request.status = RequestStatus.FINISHED_STOPPED + + # TODO: remove special handling for prompt token ids ? + if chunk_id == 0: + request.prompt_token_ids = payload_data.get("code_predictor_codes", []) + else: + request.prompt_token_ids += payload_data.get("code_predictor_codes", []) + + +def put_chunk( + connector: "OmniConnectorBase", + pooling_output: dict[str, Any], + request: Request, + custom_process_input_func: Callable[[dict[str, Any], Request], dict[str, Any] | None] | None = None, +) -> None: + """Store a chunk of pooling output. + + Args: + connector: OmniConnectorBase instance + pooling_output: Partial pooling output dictionary + request: Request object + custom_process_input_func: Optional custom function to process input + + Returns: + None: This function sends data via connector + """ + stage_id = connector.stage_id + next_stage_id = stage_id + 1 + request_id = request.request_id[0:25] + prompt_token_ids = request.prompt_token_ids + connector.request_prompt_token_ids[request_id] = prompt_token_ids + chunk_id = connector.put_requests[request_id] + connector_put_key = f"{request_id}_{stage_id}_{chunk_id}" + payload_data = None + + # TODO: add default process_input_func to handle the payload_data ? + if custom_process_input_func: + try: + payload_data = custom_process_input_func( + pooling_output=pooling_output, + request=request, + ) + except Exception as e: + logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}") + + if not payload_data: + logger.warning(f"[Stage-{stage_id}] No payload data to send for request {request_id}") + return + + if stage_id == 0 and chunk_id == 0: + if connector.request_payload.get(request_id) is None: + connector.request_payload[request_id] = payload_data + return + else: + save_payload = connector.request_payload.get(request_id) + payload_data["thinker_embeddings"] = torch.cat( + (save_payload.get("thinker_embeddings"), payload_data.get("thinker_embeddings")), dim=0 + ) + payload_data["thinker_hidden_states"] = torch.cat( + (save_payload.get("thinker_hidden_states"), payload_data.get("thinker_hidden_states")), dim=0 + ) + logger.info(f"[Stage-{stage_id}] Merged embeddings and hidden states for request {request_id}") + + if stage_id == 1: + # TODO: Make parameters configurable and optimize algorithms + chunk_size = left_context_size = 25 + connector.code_prompt_token_ids[request_id].append(payload_data.get("code_predictor_codes", [])) + length = len(connector.code_prompt_token_ids[request_id]) + chunk_length = length % chunk_size + if chunk_length != 0 and not payload_data.get("finished"): + return + + context_length = chunk_length if chunk_length != 0 else chunk_size + end_index = min(length, left_context_size + context_length) + payload_data["code_predictor_codes"] = ( + torch.tensor(connector.code_prompt_token_ids[request_id][-end_index:]) + .transpose(0, 1) + .reshape(-1) + .tolist() + ) + + success, size, metadata = connector.put( + from_stage=str(stage_id), to_stage=str(next_stage_id), put_key=connector_put_key, data=payload_data + ) + + if success: + connector.put_requests[request_id] += 1 + logger.info(f"[Stage-{stage_id}] Sent {connector_put_key}") + + +def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int: + """Compute the length of the talker prompt ids. + + Args: + prompt_ids: The prompt ids tensor. + + Returns: + The length of the talker prompt ids. + """ + im_start_token_id = 151644 + system_token_id = 8948 + user_token_id = 872 + assistant_token_id = 77091 + im_start_indexes = [i for i in range(len(prompt_ids)) if prompt_ids[i] == im_start_token_id] + im_start_indexes.append(len(prompt_ids)) + sum_user_len = 0 + assistant_len = 0 + for i in range(len(im_start_indexes) - 1): + s = im_start_indexes[i] + e = im_start_indexes[i + 1] + role = prompt_ids[s + 1] + if role == system_token_id: + continue + elif role == user_token_id: + sum_user_len += e - s + elif role == assistant_token_id and i == len(im_start_indexes) - 2: + assistant_len += 9 # 3 + 4 + 1 + 1 + else: + pass + + return sum_user_len + assistant_len diff --git a/vllm_omni/distributed/omni_connectors/connectors/base.py b/vllm_omni/distributed/omni_connectors/connectors/base.py index 6163e7b86..e3562b896 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/base.py +++ b/vllm_omni/distributed/omni_connectors/connectors/base.py @@ -13,15 +13,13 @@ class OmniConnectorBase(ABC): """Base class for all OmniConnectors.""" @abstractmethod - def put( - self, from_stage: str, to_stage: str, request_id: str, data: Any - ) -> tuple[bool, int, dict[str, Any] | None]: + def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]: """Store Python object, internal serialization handled by connector. Args: from_stage: Source stage identifier to_stage: Destination stage identifier - request_id: Unique request identifier + put_key: Unique request identifier data: Python object to store Returns: @@ -31,16 +29,13 @@ class OmniConnectorBase(ABC): pass @abstractmethod - def get( - self, from_stage: str, to_stage: str, request_id: str, metadata: dict[str, Any] | None = None - ) -> tuple[Any, int] | None: + def get(self, from_stage: str, to_stage: str, get_key: str, metadata=None) -> tuple[Any, int] | None: """Retrieve Python object and payload size (bytes). Args: from_stage: Source stage identifier to_stage: Destination stage identifier - request_id: Unique request identifier - metadata: Optional transport-specific metadata from the put operation + get_key: Unique request identifier Returns: Tuple of (Python object, serialized byte size) if found, None otherwise diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index f65e29787..a9587e21c 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time +from collections import defaultdict from typing import Any from vllm_omni.entrypoints.stage_utils import shm_read_bytes, shm_write_bytes @@ -19,6 +21,14 @@ class SharedMemoryConnector(OmniConnectorBase): def __init__(self, config: dict[str, Any]): self.config = config + self.stage_id = config.get("stage_id", -1) + self.device = config.get("device", "cuda:0") + self.put_requests: dict[str, int] = defaultdict(int) + self.get_requests: dict[str, int] = defaultdict(int) + self.finished_requests: set[str] = set() + self.request_payload = {} + self.request_prompt_token_ids: dict[str, list[int]] = defaultdict(list) + self.code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list) # Default threshold matches legacy behavior (64KB) self.threshold = int(config.get("shm_threshold_bytes", 65536)) self._metrics = { @@ -29,9 +39,7 @@ class SharedMemoryConnector(OmniConnectorBase): "inline_writes": 0, } - def put( - self, from_stage: str, to_stage: str, request_id: str, data: Any - ) -> tuple[bool, int, dict[str, Any] | None]: + def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]: try: # Always serialize first to check size (and for SHM writing) # Note: For extremely large objects in "inline" mode (e.g. Ray), @@ -40,18 +48,19 @@ class SharedMemoryConnector(OmniConnectorBase): payload = self.serialize_obj(data) size = len(payload) - if size > self.threshold: + metadata = {} + # if size > self.threshold: + if True: # TODO: correct put & get logic # Use Shared Memory - shm_name = f"omni_{from_stage}_to_{to_stage}_{request_id}" - meta = shm_write_bytes(payload, name=shm_name) + meta = shm_write_bytes(payload, name=put_key) # meta contains {'name': ..., 'size': ...} - metadata = {"shm": meta, "size": size} + metadata[put_key] = {"shm": meta, "size": size} self._metrics["shm_writes"] += 1 else: # Inline - pass bytes directly to avoid double serialization of the object # We already serialized it to check size, so we pass the bytes. # The Queue will pickle these bytes (fast), avoiding re-serializing the complex object. - metadata = {"inline_bytes": payload, "size": size} + metadata[put_key] = {"inline_bytes": payload, "size": size} self._metrics["inline_writes"] += 1 self._metrics["puts"] += 1 @@ -60,83 +69,40 @@ class SharedMemoryConnector(OmniConnectorBase): return True, size, metadata except Exception as e: - logger.error(f"SharedMemoryConnector put failed for req {request_id}: {e}") + logger.error(f"SharedMemoryConnector put failed for req {put_key}: {e}") return False, 0, None - def get( - self, from_stage: str, to_stage: str, request_id: str, metadata: dict[str, Any] | None = None - ) -> tuple[Any, int] | None: - # Helper to try reading SHM with a specific name - def try_read_shm(name: str): - # _shm is imported at module level as: import multiprocessing.shared_memory as _shm - # but to be safe inside this closure if not picked up: - from multiprocessing import shared_memory as shm_pkg + def get(self, from_stage: str, to_stage: str, get_key: str, metadata=None) -> tuple[Any, int] | None: + from multiprocessing import shared_memory as shm_pkg - shm = shm_pkg.SharedMemory(name=name) + # Wait for shared memory to be available (with retry logic) + max_retries = 30 + retry_delay = 0.1 # 100ms between retries + shm = None + + for attempt in range(max_retries): try: - # We assume the buffer size matches the data size - return shm_read_bytes({"name": name, "size": shm.size}) - finally: - shm.close() + shm = shm_pkg.SharedMemory(name=get_key) + break # Successfully opened, exit retry loop + except FileNotFoundError: + if attempt < max_retries - 1: + time.sleep(retry_delay) + else: + # Max retries reached, return None + logger.warning(f"Shared memory '{get_key}' not found after {max_retries} retries") + return None, 0 - if not metadata: - # Try to infer metadata from request_id for KV cache transfer scenarios - # where metadata is not passed out-of-band. - shm_name = f"omni_{from_stage}_to_{to_stage}_{request_id}" - - # We need the size to read from SHM using stage_utils.shm_read_bytes - try: - # Use local import to avoid potential scope issues if _shm isn't available - from multiprocessing import shared_memory as shm_pkg - - temp_shm = shm_pkg.SharedMemory(name=shm_name) - size = temp_shm.size - temp_shm.close() - - metadata = {"shm": {"name": shm_name, "size": size}, "is_fallback": True} - logger.debug(f"SharedMemoryConnector inferred metadata for {request_id}: {metadata}") - except Exception as e: - # If file not found or other error, standard path will likely fail too, but we let it try or return None - logger.debug(f"SharedMemoryConnector failed to infer metadata for {request_id} (fallback): {e}") - # We can't proceed without size if we rely on shm_read_bytes - return None + if shm is None: + return None, 0 try: - obj = None - size = 0 + data_bytes = shm_read_bytes({"name": get_key, "size": shm.size}) + obj = self.deserialize_obj(data_bytes) + return obj, shm.size + finally: + shm.close() - if "shm" in metadata: - meta = metadata["shm"] - # shm_read_bytes handles reading and unlinking - data_bytes = shm_read_bytes(meta) - obj = self.deserialize_obj(data_bytes) - size = metadata.get("size", len(data_bytes)) - elif "inline_bytes" in metadata: - # Deserialize bytes back to object - payload = metadata["inline_bytes"] - obj = self.deserialize_obj(payload) - size = metadata.get("size", len(payload)) - elif "inline" in metadata: - obj = metadata["inline"] - size = metadata.get("size", 0) - if size == 0: - # Fallback if size wasn't recorded - try: - size = len(self.serialize_obj(obj)) - except Exception: - pass - else: - logger.error( - f"Unknown metadata format in SharedMemoryConnector for req {request_id}: {list(metadata.keys())}" - ) - return None - - self._metrics["gets"] += 1 - return obj, size - - except Exception as e: - logger.error(f"SharedMemoryConnector get failed for req {request_id}: {e}") - return None + # TODO: update another read method def cleanup(self, request_id: str) -> None: # SHM segments are automatically unlinked during 'get' (shm_read_bytes). diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index 60fe62c13..91322e63e 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -91,6 +91,8 @@ def get_connectors_config_for_stage(transfer_config: OmniTransferConfig | None, # (Worker needs to create connectors to receive data) if to_stage == target_stage: stage_connectors_config[f"from_stage_{from_stage}"] = {"spec": {"name": spec.name, "extra": spec.extra}} + elif from_stage == target_stage and target_stage == "0": + stage_connectors_config[f"to_stage_{to_stage}"] = {"spec": {"name": spec.name, "extra": spec.extra}} return stage_connectors_config diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 3f2525ca8..cbe2001d7 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeTextConfig from vllm.engine.arg_utils import EngineArgs @@ -56,6 +56,11 @@ class OmniEngineArgs(EngineArgs): engine_output_type: Optional output type specification for the engine. Used to route outputs to appropriate processors (e.g., "image", "audio", "latents"). If None, output type is inferred. + custom_process_next_stage_input_func: Optional path to a custom function for processing + inputs from previous stages + If None, default processing is used. + stage_connector_spec: Extra configuration for stage connector + async_chunk: If set to True, perform async chunk """ stage_id: int = 0 @@ -63,6 +68,9 @@ class OmniEngineArgs(EngineArgs): model_arch: str = "Qwen2_5OmniForConditionalGeneration" engine_output_type: str | None = None hf_config_name: str | None = None + custom_process_next_stage_input_func: str | None = None + stage_connector_spec: dict[str, any] = field(default_factory=dict) + async_chunk: bool = False omni_kv_config: dict | None = None def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig: @@ -112,10 +120,20 @@ class OmniEngineArgs(EngineArgs): # Add the new omni-specific fields config_dict["stage_id"] = self.stage_id + config_dict["async_chunk"] = self.async_chunk config_dict["model_stage"] = self.model_stage config_dict["model_arch"] = self.model_arch config_dict["engine_output_type"] = self.engine_output_type + # Build stage_connector_config from stage_connector_spec + stage_connector_config = { + "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"), + "extra": self.stage_connector_spec.get("extra", {}).copy(), + } + stage_connector_config["extra"]["stage_id"] = self.stage_id + config_dict["stage_connector_config"] = stage_connector_config + config_dict["hf_config_name"] = self.hf_config_name + config_dict["custom_process_next_stage_input_func"] = self.custom_process_next_stage_input_func config_dict["omni_kv_config"] = self.omni_kv_config if self.hf_config_name is not None: config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict) @@ -140,6 +158,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): engine_output_type: Optional output type specification for the engine. Used to route outputs to appropriate processors (e.g., "image", "audio", "latents"). If None, output type is inferred. + stage_connector_spec: Extra configuration for stage connector """ stage_id: int = 0 @@ -147,6 +166,9 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): model_arch: str = "Qwen2_5OmniForConditionalGeneration" engine_output_type: str | None = None hf_config_name: str | None = None + custom_process_next_stage_input_func: str | None = None + stage_connector_spec: dict[str, any] = field(default_factory=dict) + async_chunk: bool = False omni_kv_config: dict | None = None def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig: @@ -186,11 +208,19 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): # Add the new omni-specific fields config_dict["stage_id"] = self.stage_id + config_dict["async_chunk"] = self.async_chunk config_dict["model_stage"] = self.model_stage config_dict["model_arch"] = self.model_arch config_dict["engine_output_type"] = self.engine_output_type + stage_connector_config = { + "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"), + "extra": self.stage_connector_spec.get("extra", {}).copy(), + } + stage_connector_config["extra"]["stage_id"] = self.stage_id + config_dict["stage_connector_config"] = stage_connector_config config_dict["hf_config_name"] = self.hf_config_name + config_dict["custom_process_next_stage_input_func"] = self.custom_process_next_stage_input_func config_dict["omni_kv_config"] = self.omni_kv_config if self.hf_config_name is not None: config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict) diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 6aef7fb16..ab5cf878c 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -111,9 +111,15 @@ class OmniRequestState(RequestState): for k, v in self.mm_accumulated.items(): if isinstance(v, list) and v and isinstance(v[0], torch.Tensor): try: - self.mm_accumulated[k] = torch.cat(v, dim=0) + if k == "audio": + # When the audio tensor shape is inconsistent, torch.cat will fail. + # We need to use torch.cat in -1 dimension. + self.mm_accumulated[k] = torch.cat(v, dim=-1) + else: + self.mm_accumulated[k] = torch.cat(v, dim=0) except Exception: # Keep last tensor on failure + logger.warning(f"Error concatenating tensor for key {k}; keeping last tensor") self.mm_accumulated[k] = v[-1] elif isinstance(v, dict): for sk, sv in v.items(): diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 3c275147f..88a290cdf 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -20,7 +20,7 @@ from vllm.v1.engine.exceptions import EngineDeadError # Internal imports (our code) from vllm_omni.config import OmniModelConfig from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector +from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length, try_send_via_connector from vllm_omni.distributed.ray_utils.utils import try_close_ray from vllm_omni.engine.input_processor import OmniInputProcessor from vllm_omni.entrypoints.client_request_state import ClientRequestState @@ -325,139 +325,76 @@ class AsyncOmni(OmniBase): self._enable_stats, _wall_start_ts, ) - # Seed stage-0 queue with all requests - logger.debug(f"[{self._name}] Seeding request into stage-0") req_state = ClientRequestState(request_id) req_state.metrics = metrics self.request_states[request_id] = req_state - # Mark first input time for stage-0 - metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time() - sp0: SamplingParams = sampling_params_list[0] # type: ignore[index] - task = { - "request_id": request_id, - "engine_inputs": prompt, - "sampling_params": sp0, - } - self.stage_list[0].submit(task) - _req_start_ts[request_id] = time.time() - logger.debug(f"[{self._name}] Enqueued request {request_id} to stage-0") + if self.async_chunk: + stage_queues = {stage_id: asyncio.Queue() for stage_id in range(num_stages)} + req_state.stage_queues = stage_queues + for i in range(num_stages): + sp: SamplingParams = sampling_params_list[i] + engine_inputs = prompt + if i != 0: + prompt_token_ids = prompt["prompt_token_ids"] + prompt_1 = prompt.copy() + prompt_1["prompt_token_ids"] = [0] * compute_talker_prompt_ids_length(prompt_token_ids) + prompt_1["multi_modal_data"] = prompt_1["mm_processor_kwargs"] = None + engine_inputs = prompt_1 - logger.debug(f"[{self._name}] Entering scheduling loop: stages={num_stages}") - for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]): - finished = False - while not finished: - result = await req_state.queue.get() - assert stage_id == req_state.stage_id + task = { + "request_id": request_id, + "engine_inputs": engine_inputs, + "sampling_params": sp, + } + self.stage_list[i].submit(task) + metrics.stage_first_ts[i] = metrics.stage_first_ts[0] or time.time() - req_id = result.get("request_id") - if "error" in result: - logger.error( - f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}", - ) - raise RuntimeError(result) # Request Finished due to error + logger.info(f"[{self._name}] Enqueued request {request_id} to stage-{str(i)}") - engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm") - if isinstance(engine_outputs, list): - engine_outputs = engine_outputs[0] - finished = engine_outputs.finished + _req_start_ts[request_id] = time.time() - # Mark last output time for this stage whenever we receive outputs - metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time()) - try: - _m = asdict(result.get("metrics")) - if _m is not None and finished: - metrics.on_stage_metrics(stage_id, req_id, _m) - except Exception as e: - logger.exception( - f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}", - ) - logger.debug( - f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing", - ) + logger.info( + f"[{self._name}] Entering scheduling loop: " + f"stages={num_stages}, final_stage={final_stage_id_for_e2e}" + ) + async for output in self._process_async_results( + request_id, + req_state, + metrics, + final_stage_id_for_e2e, + _req_start_ts, + _wall_start_ts, + ): + yield output + else: + sp0: SamplingParams = sampling_params_list[0] # type: ignore[index] + task = { + "request_id": request_id, + "engine_inputs": prompt, + "sampling_params": sp0, + } + self.stage_list[0].submit(task) - if getattr(stage, "final_output", False): - logger.debug( - f"[{self._name}] Request {req_id} finalized at stage-{stage_id}", - ) + _req_start_ts[request_id] = time.time() + # Mark first input time for stage-0 + metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time() + logger.info( + f"[{self._name}] Entering scheduling loop: " + f"stages={num_stages}, final_stage={final_stage_id_for_e2e}" + ) - # End-to-end timing and time-per-token for final output - # (only once per request at the designated final stage) - try: - rid_key = str(req_id) - if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done and finished: - metrics.on_finalize_request( - stage_id, - req_id, - _req_start_ts.get(req_id, _wall_start_ts), - ) - except Exception as e: - logger.exception( - f"[{self._name}] Finalize request handling error for req " - f"{req_id} at stage {stage_id}: {e}", - ) - - # Handle diffusion outputs that already contain images - if stage.final_output_type == "image": - images = [] - if isinstance(engine_outputs, OmniRequestOutput) and engine_outputs.images: - images = engine_outputs.images - elif hasattr(engine_outputs, "images") and engine_outputs.images: - images = engine_outputs.images - yield OmniRequestOutput( - stage_id=stage_id, - final_output_type=stage.final_output_type, - request_output=engine_outputs, - images=images, - ) - else: - yield OmniRequestOutput( - stage_id=stage_id, - final_output_type=stage.final_output_type, - request_output=engine_outputs, - ) - if not isinstance(engine_outputs, list): - engine_outputs = [engine_outputs] - stage.set_engine_outputs(engine_outputs) - # Forward to next stage if there is one - next_stage_id = stage_id + 1 - if next_stage_id <= final_stage_id_for_e2e and finished: - next_stage: OmniStage = self.stage_list[next_stage_id] - next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt) - sp_next: SamplingParams = sampling_params_list[next_stage_id] - - # Check if we have a connector for this edge - connector_key = (str(stage_id), str(next_stage_id)) - connector = self.connectors.get(connector_key) - - sent_via_connector = False - if connector: - sent_via_connector = try_send_via_connector( - connector=connector, - stage_id=stage_id, - next_stage_id=next_stage_id, - req_id=req_id, - next_inputs=next_inputs, - sampling_params=sp_next, - original_prompt=prompt, - next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit, - metrics=metrics, - ) - - if not sent_via_connector: - # Fallback logic removed as we now enforce connector usage. - # If no connector is found or send fails, we log an error and raise, - # because continuing would cause the request to be silently dropped - # and the orchestrator to hang waiting for completion. - error_msg = ( - f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. " - "Configure a connector for this edge or inspect connector logs for details." - ) - logger.error(error_msg) - raise RuntimeError(error_msg) - logger.debug(f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}") - else: - logger.debug(f"[{self._name}] Request {req_id} fully completed") + async for output in self._process_sequential_results( + request_id, + req_state, + metrics, + final_stage_id_for_e2e, + _req_start_ts, + _wall_start_ts, + sampling_params_list, + prompt, + ): + yield output logger.debug(f"[{self._name}] All requests completed") @@ -474,6 +411,184 @@ class AsyncOmni(OmniBase): logger.info("[AsyncOrchestrator] Request %s aborted.", request_id) raise + async def _process_async_results( + self, + request_id: str, + req_state: ClientRequestState, + metrics: OrchestratorMetrics, + final_stage_id_for_e2e: int, + req_start_ts: dict[int, float], + wall_start_ts: float, + ) -> AsyncGenerator[OmniRequestOutput, None]: + all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)} + while not all(all_stages_finished.values()): + for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]): + if all_stages_finished[stage_id]: + continue + result = await req_state.stage_queues[stage_id].get() + logger.info(f"[{self._name}] Received result from stage-{stage_id}: {result}") + engine_outputs, finished, output_to_yield = self._process_single_result( + result, stage, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e + ) + + all_stages_finished[stage_id] = finished + + if output_to_yield: + yield output_to_yield + + async def _process_sequential_results( + self, + request_id: str, + req_state: ClientRequestState, + metrics: OrchestratorMetrics, + final_stage_id_for_e2e: int, + req_start_ts: dict[int, float], + wall_start_ts: float, + sampling_params_list: list[SamplingParams], + prompt: Any, + ) -> AsyncGenerator[OmniRequestOutput, None]: + for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]): + finished = False + while not finished: + result = await req_state.queue.get() + assert stage_id == req_state.stage_id + req_id = result.get("request_id") + engine_outputs, finished, output_to_yield = self._process_single_result( + result, stage, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e + ) + if output_to_yield: + yield output_to_yield + if not isinstance(engine_outputs, list): + engine_outputs = [engine_outputs] + stage.set_engine_outputs(engine_outputs) + # Forward to next stage if there is one + next_stage_id = stage_id + 1 + if next_stage_id <= final_stage_id_for_e2e and finished: + next_stage: OmniStage = self.stage_list[next_stage_id] + next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt) + sp_next: SamplingParams = sampling_params_list[next_stage_id] + + # Check if we have a connector for this edge + connector_key = (str(stage_id), str(next_stage_id)) + connector = self.connectors.get(connector_key) + + sent_via_connector = False + if connector: + sent_via_connector = try_send_via_connector( + connector=connector, + stage_id=stage_id, + next_stage_id=next_stage_id, + req_id=req_id, + next_inputs=next_inputs, + sampling_params=sp_next, + original_prompt=prompt, + next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit, + metrics=metrics, + ) + + if not sent_via_connector: + # Fallback logic removed as we now enforce connector usage. + # If no connector is found or send fails, we log an error and raise, + # because continuing would cause the request to be silently dropped + # and the orchestrator to hang waiting for completion. + error_msg = ( + f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. " + "Configure a connector for this edge or inspect connector logs for details." + ) + logger.error(error_msg) + raise RuntimeError(error_msg) + logger.debug(f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}") + else: + logger.debug(f"[{self._name}] Request {req_id} fully completed") + + def _process_single_result( + self, + result: dict[str, Any], + stage: OmniStage, + stage_id: int, + metrics: OrchestratorMetrics, + req_start_ts: dict[int, float], + wall_start_ts: float, + final_stage_id_for_e2e: int, + ) -> tuple[Any, bool, OmniRequestOutput | None]: + """ + Process a single result dictionary from a stage. + Returns: + engine_outputs: The decoded outputs. + finished: Whether the stage processing is finished for this request. + output_to_yield: An OmniRequestOutput to yield, or None. + """ + req_id = result.get("request_id") + if "error" in result: + logger.error( + f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}", + ) + raise RuntimeError(result) + + engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm") + if isinstance(engine_outputs, list): + engine_outputs = engine_outputs[0] + + finished = engine_outputs.finished + + # Mark last output time + metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time()) + + try: + _m = asdict(result.get("metrics")) + if _m is not None and finished: + metrics.on_stage_metrics(stage_id, req_id, _m) + except Exception as e: + logger.exception( + f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}", + ) + + logger.debug( + f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing", + ) + + output_to_yield = None + if getattr(stage, "final_output", False): + logger.debug(f"[{self._name}] Request {req_id} finalized at stage-{stage_id}") + + # Finalize request metrics if this is the E2E final stage and it's finished + try: + rid_key = str(req_id) + if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done and finished: + metrics.on_finalize_request( + stage_id, + req_id, + req_start_ts.get(req_id, wall_start_ts), + ) + except Exception as e: + logger.exception( + f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}", + ) + + # Construct output to yield + images = [] + if stage.final_output_type == "image": + if isinstance(engine_outputs, OmniRequestOutput) and engine_outputs.images: + images = engine_outputs.images + elif hasattr(engine_outputs, "images") and engine_outputs.images: + images = engine_outputs.images + + if stage.final_output_type == "image": + output_to_yield = OmniRequestOutput( + stage_id=stage_id, + final_output_type=stage.final_output_type, + request_output=engine_outputs, + images=images, + ) + else: + output_to_yield = OmniRequestOutput( + stage_id=stage_id, + final_output_type=stage.final_output_type, + request_output=engine_outputs, + ) + + return engine_outputs, finished, output_to_yield + def _run_output_handler(self) -> None: if self.output_handler is not None: return @@ -503,8 +618,12 @@ class AsyncOmni(OmniBase): dropping output for req {req_id} at stage-{stage_id}" ) continue - await req_state.queue.put(result) - req_state.stage_id = stage_id + if hasattr(req_state, "stage_queues") and stage_id in req_state.stage_queues: + await req_state.stage_queues[stage_id].put(result) + else: + # Fallback to old behavior for compatibility + await req_state.queue.put(result) + req_state.stage_id = stage_id if idle: await asyncio.sleep(0.001) # Avoid CPU overload when idle else: @@ -512,7 +631,14 @@ class AsyncOmni(OmniBase): except Exception as e: logger.exception("AsyncOmni output_handler failed.") for req_state in request_states.values(): - await req_state.queue.put({"request_id": req_id, "error": str(e)}) + error_msg = {"request_id": req_state.request_id, "error": str(e)} + # Send error to all stage queues + if hasattr(req_state, "stage_queues"): + for queue in req_state.stage_queues.values(): + await queue.put(error_msg) + else: + await req_state.queue.put(error_msg) + error_msg = {"request_id": req_state.request_id, "error": str(e)} self.output_handler = None # Make possible for restart self.output_handler = asyncio.create_task(output_handler()) diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 3d918aea7..e76fe834a 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -232,6 +232,8 @@ class OmniBase: self.worker_backend = worker_backend self.ray_address = ray_address self.batch_timeout = batch_timeout + # async chunk remains the same for each stage + self.async_chunk = self._is_async_chunk_enable(self.stage_configs) # Build OmniStage instances in parallel, preserve original order def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]: @@ -261,6 +263,11 @@ class OmniBase: # Wait for all stages to report readiness before seeding self._wait_for_stages_ready(timeout=init_timeout) + def _is_async_chunk_enable(self, stage_args: list) -> bool: + """get async chunk flag""" + engine_args = getattr(stage_args[0], "engine_args", None) + return bool(getattr(engine_args, "async_chunk", False)) + def _start_stages(self, model: str) -> None: """Start all stage processes.""" if self.worker_backend == "ray": diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 7beef030d..e0e699c9c 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -673,12 +673,15 @@ def _stage_worker( e, ) # Init engine based on stage_type - logger.debug( - "[Stage-%s] Initializing %s engine with args keys=%s", - stage_id, - stage_type, - list(engine_args.keys()), - ) + logger.debug("[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys())) + if engine_args.get("async_chunk", False): + logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id) + stage_connector_spec = {} + for v in connectors_config.values(): + stage_connector_spec = dict(v.get("spec", {})) + break + engine_args["stage_connector_spec"] = stage_connector_spec + engine_args["stage_id"] = stage_id try: if stage_type == "diffusion": engine_args.pop("model_stage", None) @@ -1206,6 +1209,14 @@ async def _stage_worker_async( stage_type, list(engine_args.keys()), ) + if engine_args.get("async_chunk", False): + logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id) + stage_connector_spec = {} + for v in connectors_config.values(): + stage_connector_spec = dict(v.get("spec", {})) + break + engine_args["stage_connector_spec"] = stage_connector_spec + engine_args["stage_id"] = stage_id try: if stage_type == "diffusion": # For diffusion, we need to extract diffusion-specific config diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 8a62fd16c..5d9f76d70 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -18,6 +18,7 @@ try: except ImportError: soundfile = None + from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, @@ -1213,6 +1214,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): yield f"data: {data}\n\n" elif final_output_type == "audio": + role = self.get_chat_request_role(request) choices_data = self._create_audio_choice(omni_res, role, request, stream=True) chunk = OmniChatCompletionStreamResponse( id=request_id, @@ -1655,7 +1657,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ): choices: list[ChatCompletionResponseChoice] = [] final_res = omni_outputs.request_output - audio_tensor = final_res.multimodal_output["audio"].float().detach().cpu().numpy() + if stream: + audio_tensor = final_res.multimodal_output["audio"][-1].float().detach().cpu().numpy() + else: + audio_tensor = final_res.multimodal_output["audio"].float().detach().cpu().numpy() # Ensure audio is 1D (flatten if needed) if audio_tensor.ndim > 1: diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index db17116ae..da837ce9a 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -222,7 +222,8 @@ def maybe_dump_to_shm(obj: Any, threshold: int) -> tuple[bool, Any]: """ payload = serialize_obj(obj) if len(payload) > threshold: - return True, shm_write_bytes(payload) + logger.debug(f"Dumping object to SHM with size: {len(payload)}") + return True, shm_write_bytes(payload, name=None) return False, obj diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index bbb75a968..f1ed4d843 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -221,6 +221,7 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None base_engine_args = {} config_data = OmegaConf.load(config_path) stage_args = config_data.stage_args + global_async_chunk = config_data.get("async_chunk", False) # Convert any nested dataclass objects to dicts before creating OmegaConf base_engine_args = _convert_dataclasses_to_dict(base_engine_args) base_engine_args = OmegaConf.create(base_engine_args) @@ -234,6 +235,7 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None runtime_cfg = stage_arg.runtime max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1) base_engine_args_tmp["max_num_seqs"] = max_batch_size + base_engine_args_tmp.async_chunk = global_async_chunk stage_arg.engine_args = base_engine_args_tmp return stage_args diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 17f16d1d9..0dd590710 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -180,6 +180,8 @@ class Qwen3OmniMoeForConditionalGeneration( self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None ) + self.chunk_segment_info = {} # request_id -> chunk_segment_info + # ==================== Device utilities ==================== @staticmethod @@ -483,12 +485,19 @@ class Qwen3OmniMoeForConditionalGeneration( talker_codes = talker_codes.unsqueeze(0).unsqueeze(0) talker_codes = talker_codes.expand(1, 16, -1) - # Use chunked decode for memory efficiency - audio_tensor = self.code2wav.chunked_decode( - talker_codes, - chunk_size=300, - left_context_size=25, - ) + if self.vllm_config.model_config.async_chunk: + audio_tensor = self.code2wav.chunked_decode_streaming( + talker_codes, + chunk_size=25, + left_context_size=25, + ) + else: + # Use chunked decode for memory efficiency + audio_tensor = self.code2wav.chunked_decode( + talker_codes, + chunk_size=300, + left_context_size=25, + ) return audio_tensor @@ -595,7 +604,6 @@ class Qwen3OmniMoeForConditionalGeneration( last_talker_hidden, text_step, update_dict = self.talker_preprocess_decode( input_ids, input_embeds, **info_dict ) - update_dict["mtp_inputs"] = last_talker_hidden, text_step return input_ids, input_embeds, update_dict @@ -624,6 +632,33 @@ class Qwen3OmniMoeForConditionalGeneration( inputs_embeds = (inputs_embeds + text_step).reshape(-1, self.talker_config.text_config.hidden_size) return inputs_embeds, code_predictor_codes.squeeze(-1) + def _get_tts_embed(self, thinker_embed, tts_bos_thinker, tts_eos_thinker, tts_pad_thinker): + """Project thinker-side TTS embeddings into talker text space.""" + module_device = self._module_device(self.talker) + + def _ensure_1x1(x: torch.Tensor) -> torch.Tensor: + if x.ndim == 3: + return x[0, -1:, :] + if x.ndim == 2: + return x[-1] + return x.view(1, 1, -1) + + def _proj_from_thinker(x_opt: torch.Tensor | None) -> torch.Tensor: + if isinstance(x_opt, torch.Tensor) and x_opt.numel() > 0: + xin = _ensure_1x1(x_opt).to(module_device) + else: + xin = torch.zeros( + (1, thinker_embed.shape[-1]), + device=module_device, + dtype=thinker_embed.dtype, + ) + return self.talker.text_projection(xin).to(module_device) + + self.tts_bos_embed = _proj_from_thinker(tts_bos_thinker) + self.tts_eos_embed = _proj_from_thinker(tts_eos_thinker) + self.tts_pad_embed = _proj_from_thinker(tts_pad_thinker) + return self.tts_bos_embed, self.tts_eos_embed, self.tts_pad_embed + def talker_preprocess_prefill(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict): # Containers to return per-request updates (e.g., code_predictor_hidden_per_request) update_dict: dict[str, dict] = {} @@ -647,7 +682,7 @@ class Qwen3OmniMoeForConditionalGeneration( if info_dict.get("thinker_input_ids") is None else torch.as_tensor(info_dict.get("thinker_input_ids"), device=self._module_device(self.talker)) ) - # Optional thinker-side TTS token embeddings (to be projected by talker) + tts_bos_thinker = info_dict.get("tts_bos_embed").to( device=self._module_device(self.talker), dtype=torch.bfloat16 ) @@ -700,7 +735,7 @@ class Qwen3OmniMoeForConditionalGeneration( tts_pad_thinker=tts_pad_thinker, ) - # Queue tailing_text_hidden for decode (drop first for next steps), + # Queue trailing_text_hidden for decode (drop first for next steps), try: if isinstance(trailing_text_hidden, torch.Tensor) and trailing_text_hidden.numel() > 0: if trailing_text_hidden.ndim == 2: @@ -716,7 +751,7 @@ class Qwen3OmniMoeForConditionalGeneration( # compatible with old shape [1,S,D] rem_tail = trailing_text_hidden.squeeze(0) if rem_tail.shape[0] > 0: - update_dict["tailing_text_hidden"] = rem_tail.detach().to("cpu").contiguous() + update_dict["trailing_text_hidden"] = rem_tail.detach().to("cpu").contiguous() # Also persist projected tts_pad for decode fallback if needed if isinstance(tts_pad_thinker, torch.Tensor): pad_in = tts_pad_thinker @@ -762,31 +797,9 @@ class Qwen3OmniMoeForConditionalGeneration( (thinker_result_ids == self.thinker_config.video_token_id) ).to(input_ids.device) # [t] # fmt: skip - # Project thinker-side TTS embeddings into talker text space (provided by thinker stage) - talker_dev = self._module_device(self.talker) - - def _ensure_1x1(x: torch.Tensor) -> torch.Tensor: - if x.ndim == 3: - return x[0, -1:, :] - if x.ndim == 2: - return x[-1] - return x.view(1, 1, -1) - - def _proj_from_thinker(x_opt: torch.Tensor | None) -> torch.Tensor: - if isinstance(x_opt, torch.Tensor) and x_opt.numel() > 0: - xin = _ensure_1x1(x_opt).to(talker_dev) - else: - xin = torch.zeros( - (1, thinker_embed.shape[-1]), - device=talker_dev, - dtype=thinker_embed.dtype, - ) - return self.talker.text_projection(xin).to(input_ids.device) - - tts_bos_embed = _proj_from_thinker(tts_bos_thinker) - tts_eos_embed = _proj_from_thinker(tts_eos_thinker) - tts_pad_embed = _proj_from_thinker(tts_pad_thinker) - self.tts_pad_embed = tts_pad_embed + tts_bos_embed, tts_eos_embed, tts_pad_embed = self._get_tts_embed( + thinker_embed, tts_bos_thinker, tts_eos_thinker, tts_pad_thinker + ) talker_input_embeds = [] # [1 t d] talker_input_ids = [] @@ -835,24 +848,58 @@ class Qwen3OmniMoeForConditionalGeneration( return talker_input_id, talker_input_embed, trailing_text_hidden_all + def _thinker_decode_to_talker_decode( + self, + info_dict: dict, + device: torch.device, + update_dict, + ): + """ + Project thinker outputs to talker inputs during prefill stage. + Returns: + (input_ids, input_embeds) for talker + """ + thinker_embed = info_dict.get("thinker_embeddings", None) + if thinker_embed is None: + if info_dict.get("finished_flag"): + return self.tts_pad_embed.to(device) + update_dict["finished_flag"] = True + return self.tts_eos_embed.to(device) + + thinker_embed = thinker_embed.to(device) + return self.talker.text_projection(thinker_embed).to(device) + def talker_preprocess_decode(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict): update_dict: dict[str, dict] = {} + last_talker_hidden = None + text_step = None try: - q_tail = info_dict.get("tailing_text_hidden") - if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0: - use_vec = q_tail[0:1, :] - new_q_tail = ( - q_tail[1:, :].detach().to("cpu").contiguous() - if q_tail.shape[0] > 1 - else self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype) - ) - text_step = use_vec.to(input_embeds.device, dtype=input_embeds.dtype) - update_dict["tailing_text_hidden"] = new_q_tail + if self.vllm_config.model_config.async_chunk: + text_step = self._thinker_decode_to_talker_decode(info_dict, input_ids.device, update_dict) else: - text_step = self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype) + q_tail = info_dict.get("trailing_text_hidden", None) + if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0: + use_vec = q_tail[0:1, :] + new_q_tail = ( + q_tail[1:, :].detach().to("cpu").contiguous() + if q_tail.shape[0] > 1 + else self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype) + ) + text_step = use_vec.to(input_embeds.device, dtype=input_embeds.dtype) + update_dict["trailing_text_hidden"] = new_q_tail + else: + text_step = self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype) - last_talker_hidden = info_dict.get("last_talker_hidden").to(input_embeds.device, dtype=input_embeds.dtype) - last_talker_hidden = last_talker_hidden.reshape(*last_talker_hidden.shape[-2:]) # [1, hidden_size] + last_talker_hidden_tensor = info_dict.get("last_talker_hidden") + if last_talker_hidden_tensor is not None: + last_talker_hidden = last_talker_hidden_tensor.to(input_embeds.device, dtype=input_embeds.dtype) + last_talker_hidden = last_talker_hidden.reshape(*last_talker_hidden.shape[-2:]) # [1, hidden_size] + else: + last_talker_hidden = torch.zeros( + (1, self.talker_config.text_config.hidden_size), + device=input_embeds.device, + dtype=input_embeds.dtype, + ) except Exception as e: logger.error(f"Error in decode: {e}") @@ -866,7 +913,6 @@ class Qwen3OmniMoeForConditionalGeneration( ) user_mm_mask = multimodal_mask[im_start_index:segment_end_index] - # Multimodal data exists if user_mm_mask.any(): user_thinker_hidden_mm = thinker_hidden[im_start_index:segment_end_index][user_mm_mask] @@ -883,12 +929,20 @@ class Qwen3OmniMoeForConditionalGeneration( assistant_hidden = self.talker.text_projection(thinker_embed[im_start_index:segment_end_index]).to( tts_pad_embed.device ) # [t, d] + + # [3 tokens] + [4 pad] + [1 BOS] + [1 first text] = 9 tokens assistant_text_hidden = torch.cat( ( assistant_hidden[:3], tts_pad_embed.expand(4, -1), tts_bos_embed, - assistant_hidden[3:4], # First text + assistant_hidden[3:4] + if assistant_hidden.shape[0] > 3 + else torch.zeros( + (1, assistant_hidden.shape[1]), + device=assistant_hidden.device, + dtype=assistant_hidden.dtype, + ), # First text ), dim=0, ) @@ -904,6 +958,9 @@ class Qwen3OmniMoeForConditionalGeneration( device=tts_pad_embed.device, dtype=torch.long, ) + embed_input_ids = self.talker.embed_input_ids(codec_special_tokens).to( + device=tts_pad_embed.device, dtype=torch.bfloat16 + ) assistant_codec_hidden = torch.cat( ( torch.zeros( @@ -911,18 +968,21 @@ class Qwen3OmniMoeForConditionalGeneration( device=tts_pad_embed.device, dtype=torch.bfloat16, ), - self.talker.embed_input_ids(codec_special_tokens).to(device=tts_pad_embed.device, dtype=torch.bfloat16), - ), - dim=0, - ) - trailing_text_hidden = torch.cat( - ( - assistant_hidden[4:], - tts_eos_embed, + embed_input_ids, ), dim=0, ) + if assistant_hidden.shape[0] > 4: + trailing_text_hidden = torch.cat( + (assistant_hidden[4:], tts_eos_embed), + dim=0, + ) + else: + trailing_text_hidden = torch.zeros( + tts_eos_embed.shape, device=tts_eos_embed.device, dtype=tts_eos_embed.dtype + ) + input_embeds = assistant_text_hidden + assistant_codec_hidden input_ids = torch.full( (assistant_text_hidden.shape[0],), diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index f430ecb16..bc2242f70 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -196,6 +196,39 @@ class Qwen3OmniMoeCode2Wav(nn.Module): return torch.cat(wavs, dim=-1) + def chunked_decode_streaming( + self, + codes: torch.Tensor, + chunk_size: int = 25, + left_context_size: int = 25, + ) -> torch.Tensor: + """ + Decode long sequences in chunks to avoid OOM. + + Uses overlapping chunks with left context to avoid boundary artifacts. + + Args: + codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes + chunk_size: Number of codec frames per chunk + left_context_size: Number of overlapping frames for context + + Returns: + waveform: [batch, 1, waveform_len] - Complete waveform + """ + wavs = [] + end_index = codes.shape[-1] + # TODO: need to optimize algorithms, current only support + # chunk_size = left_context_size = 25 + if end_index == chunk_size: + context_size = 0 + else: + context_size = left_context_size + # Decode chunk + wav_chunk = self(codes) + # Remove context from output (context_size * total_upsample samples) + wavs.append(wav_chunk[..., context_size * self.total_upsample :]) + return torch.cat(wavs, dim=-1) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights from HuggingFace checkpoint.""" loader = AutoWeightsLoader( diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml index d4de07823..9c422d76e 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml @@ -4,12 +4,13 @@ # Stage 2: Code2Wav (8-layer RVQ codes → audio waveform) # The following config has been verified on 2x H100-80G GPUs. +async_chunk: false stage_args: - stage_id: 0 stage_type: llm # Use llm stage type to launch OmniLLM runtime: devices: "0,1" - max_batch_size: 1 + max_batch_size: 64 engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration @@ -40,17 +41,16 @@ stage_args: stage_type: llm # Use llm stage type to launch OmniLLM runtime: devices: "1" - max_batch_size: 1 + max_batch_size: 64 engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.35 + gpu_memory_utilization: 0.3 enforce_eager: true trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav - # tensor_parallel_size: 2 enable_prefix_caching: false max_num_batched_tokens: 32768 distributed_executor_backend: "mp" diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml new file mode 100644 index 000000000..55d455f56 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml @@ -0,0 +1,101 @@ +# Stage config for running Qwen3-Omni-MoE with 3-stage architecture +# Stage 0: Thinker (multimodal understanding + text generation) +# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes) +# Stage 2: Code2Wav (16-layer RVQ codes → audio waveform) + +# The following config has been verified on 2x H100-80G GPUs. +async_chunk: true +stage_args: + - stage_id: 0 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + devices: "0,1" + max_batch_size: 64 + engine_args: + model_stage: thinker + model_arch: Qwen3OmniMoeForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.6 + enforce_eager: false + trust_remote_code: true + engine_output_type: latent # Output hidden states for talker + distributed_executor_backend: "mp" + enable_prefix_caching: false + max_num_batched_tokens: 32768 + hf_config_name: thinker_config + tensor_parallel_size: 2 + custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk + final_output: true + final_output_type: text + is_comprehension: true + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + top_k: 1 + max_tokens: 2048 + seed: 42 + detokenize: True + repetition_penalty: 1.05 + + - stage_id: 1 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + devices: "0" + max_batch_size: 64 + engine_args: + model_stage: talker + model_arch: Qwen3OmniMoeForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.3 + enforce_eager: false + trust_remote_code: true + engine_output_type: latent # Output codec codes for code2wav + enable_prefix_caching: false + max_num_batched_tokens: 32768 + distributed_executor_backend: "mp" + hf_config_name: talker_config + custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk + engine_input_source: [0] + # final_output: true + # final_output_type: text + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + seed: 42 + detokenize: False + repetition_penalty: 1.05 + stop_token_ids: [2150] + + - stage_id: 2 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + devices: "1" + max_batch_size: 1 + engine_args: + model_stage: code2wav + model_arch: Qwen3OmniMoeForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio # Final output: audio waveform + gpu_memory_utilization: 0.1 + distributed_executor_backend: "mp" + max_num_batched_tokens: 10000 + hf_config_name: thinker_config + engine_input_source: [1] + final_output: true + final_output_type: audio + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + detokenize: True + repetition_penalty: 1.1 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index a1457a975..03daa8e42 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -9,6 +9,7 @@ import torch from vllm.inputs import TextPrompt from vllm.platforms import current_platform +from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt @@ -48,6 +49,74 @@ def _compute_talker_prompt_ids_length(info, device: torch.device | str = "cuda") return sum_user_len + assistant_len +# ========================= +# Common helpers +# ========================= + + +def _ensure_list(x): + """Convert ConstantList / tensor-like to Python list.""" + if hasattr(x, "_x"): + return list(x._x) + elif not isinstance(x, list): + return x + return list(x) + + +def _validate_stage_inputs(stage_list, engine_input_source): + if not engine_input_source: + raise ValueError("engine_input_source cannot be empty") + + stage_id = engine_input_source[0] + if stage_id >= len(stage_list): + raise IndexError(f"Invalid stage_id: {stage_id}") + + stage = stage_list[stage_id] + if stage.engine_outputs is None: + raise RuntimeError(f"Stage {stage_id} has no outputs yet") + + return stage.engine_outputs + + +# ========================= +# Thinker -> Talker +# ========================= + + +def thinker2talker_async_chunk( + pooling_output: dict[str, Any], + request: OmniEngineCoreRequest, +) -> list[dict[str, Any]]: + """ + Process thinker outputs to create talker inputs. + 1. thinker's text generation outputs (token IDs + hidden states) + 2. Split hidden states into: prompt embeddings + generated embeddings + 3. Package for talker with additional information + """ + all_token_ids = request.all_token_ids # prefill + decode + prompt_token_ids = request.prompt_token_ids + + # Convert ConstantList to regular list for OmniSerializer serialization + all_token_ids = _ensure_list(all_token_ids) + prompt_token_ids = _ensure_list(prompt_token_ids) + + thinker_output = pooling_output + + talker_additional_info = { + "thinker_embeddings": thinker_output.get("0").detach().cpu(), + "thinker_hidden_states": thinker_output.get("24").detach().cpu(), + "thinker_sequences": all_token_ids, + "thinker_input_ids": prompt_token_ids, + # Provide thinker-side TTS token embeddings for talker projection + "tts_bos_embed": thinker_output.get("tts_bos_embed").detach().cpu(), + "tts_eos_embed": thinker_output.get("tts_eos_embed").detach().cpu(), + "tts_pad_embed": thinker_output.get("tts_pad_embed").detach().cpu(), + "finished": torch.tensor(request.is_finished(), dtype=torch.bool), + } + + return talker_additional_info + + def thinker2talker( stage_list: list[Any], engine_input_source: list[int], @@ -71,47 +140,33 @@ def thinker2talker( Returns: List of OmniTokensPrompt for talker stage """ - if not engine_input_source: - raise ValueError("engine_input_source cannot be empty") - - source_stage_id = engine_input_source[0] - if source_stage_id >= len(stage_list): - raise IndexError(f"Invalid stage_id: {source_stage_id}") - - if stage_list[source_stage_id].engine_outputs is None: - raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") - - thinker_outputs = stage_list[source_stage_id].engine_outputs - talker_inputs = [] + thinker_outputs = _validate_stage_inputs(stage_list, engine_input_source) + talker_inputs: list[OmniTokensPrompt] = [] device = torch.device(current_platform.device_type) # Process each thinker output - for i, thinker_output in enumerate(thinker_outputs): + for thinker_output in thinker_outputs: output = thinker_output.outputs[0] - thinker_embeddings = output.multimodal_output["0"].detach().to(device=device, dtype=torch.float) - thinker_hidden_states = output.multimodal_output["24"].detach().to(device=device, dtype=torch.float) info = { - "thinker_embeddings": thinker_embeddings, - "thinker_hidden_states": thinker_hidden_states, - "thinker_sequences": thinker_output.prompt_token_ids - + output.token_ids, # the thinker_sequences is the whole ids + "thinker_embeddings": output.multimodal_output["0"].detach().to(device=device, dtype=torch.float), + "thinker_hidden_states": output.multimodal_output["24"].detach().to(device=device, dtype=torch.float), + "thinker_sequences": ( + thinker_output.prompt_token_ids + output.token_ids + ), # the thinker_sequences is the whole ids "thinker_input_ids": thinker_output.prompt_token_ids, # Provide thinker-side TTS token embeddings for talker projection - "tts_bos_embed": ( - output.multimodal_output.get("tts_bos_embed").detach().to(device=device, dtype=torch.float) - ), - "tts_eos_embed": ( - output.multimodal_output.get("tts_eos_embed").detach().to(device=device, dtype=torch.float) - ), - "tts_pad_embed": ( - output.multimodal_output.get("tts_pad_embed").detach().to(device=device, dtype=torch.float) - ), + "tts_bos_embed": output.multimodal_output["tts_bos_embed"].detach().to(device=device, dtype=torch.float), + "tts_eos_embed": output.multimodal_output["tts_eos_embed"].detach().to(device=device, dtype=torch.float), + "tts_pad_embed": output.multimodal_output["tts_pad_embed"].detach().to(device=device, dtype=torch.float), } + + prompt_len = _compute_talker_prompt_ids_length(info, device=device) + talker_inputs.append( OmniTokensPrompt( - prompt_token_ids=[0] * _compute_talker_prompt_ids_length(info, device=device), + prompt_token_ids=[0] * prompt_len, additional_information=info, multi_modal_data=None, mm_processor_kwargs=None, @@ -121,6 +176,50 @@ def thinker2talker( return talker_inputs +# ========================= +# Talker -> Code2Wav +# ========================= + + +def talker2code2wav_async_chunk( + pooling_output: dict[str, Any], + request: OmniEngineCoreRequest, +): + """ + Pooling version. + """ + if "code_predictor_codes" not in pooling_output: + return [] + + code_predictor_codes = pooling_output["code_predictor_codes"] + + if code_predictor_codes is None: + return [] + if isinstance(code_predictor_codes, torch.Tensor): + if code_predictor_codes.numel() == 0: + return [] + elif hasattr(code_predictor_codes, "__len__"): + if len(code_predictor_codes) == 0: + return [] + + if isinstance(code_predictor_codes, torch.Tensor): + if not code_predictor_codes.any(): + return [] + else: + code_tensor = torch.tensor(code_predictor_codes, dtype=torch.long) + if not code_tensor.any(): + return [] + + codec_codes = code_predictor_codes.to(torch.long).transpose(0, 1).cpu().to(torch.long).reshape(-1).tolist() + if sum(codec_codes) == 0: + return [] + + return { + "code_predictor_codes": codec_codes, + "finished": torch.tensor(request.is_finished(), dtype=torch.bool), + } + + def talker2code2wav( stage_list: list[Any], engine_input_source: list[int], @@ -144,21 +243,10 @@ def talker2code2wav( Returns: List of OmniTokensPrompt for code2wav stage """ - if not engine_input_source: - raise ValueError("engine_input_source cannot be empty") - - source_stage_id = engine_input_source[0] - if source_stage_id >= len(stage_list): - raise IndexError(f"Invalid stage_id: {source_stage_id}") - - if stage_list[source_stage_id].engine_outputs is None: - raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") - - talker_outputs = stage_list[source_stage_id].engine_outputs - code2wav_inputs = [] - + talker_outputs = _validate_stage_inputs(stage_list, engine_input_source) + code2wav_inputs: list[OmniTokensPrompt] = [] # Process each talker output - for i, talker_output in enumerate(talker_outputs): + for talker_output in talker_outputs: output = talker_output.outputs[0] seq_len = len(output.token_ids) - 1 # Extract codec codes from talker output diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 6055c93f7..ecc839fd3 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -280,7 +280,6 @@ class GPUARModelRunner(OmniGPUModelRunner): aux_hidden_states = None hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output) - if multimodal_outputs is not None: keys_or_type = ( list(multimodal_outputs.keys()) @@ -501,7 +500,9 @@ class GPUARModelRunner(OmniGPUModelRunner): dtype=np.int32, ) - self._process_additional_information_updates(hidden_states, multimodal_outputs, num_scheduled_tokens_np) + self._process_additional_information_updates( + hidden_states, multimodal_outputs, num_scheduled_tokens_np, scheduler_output + ) pooler_output: list[dict[str, object]] = [] for rid in req_ids_output_copy: @@ -761,8 +762,9 @@ class GPUARModelRunner(OmniGPUModelRunner): for attempt in range(max_retries): try: + put_key = f"omni_{from_stage}_to_{to_stage}_{request_id}" success, size, metadata = connector.put( - from_stage=from_stage, to_stage=to_stage, request_id=request_id, data=data + from_stage=from_stage, to_stage=to_stage, put_key=put_key, data=data ) if success: return success, size, metadata diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 7d5802483..87e1419c8 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -50,6 +50,17 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): - Executes generation process and returns tensors via `pooler_output`. """ + def _update_request_states(self, scheduler_output: SchedulerOutput): + cached_reqs = scheduler_output.scheduled_cached_reqs + for _, req_id in enumerate(cached_reqs.req_ids): + req_state = self.requests.get(req_id) + assert req_state is not None + req_state.prompt_token_ids = cached_reqs.prompt_token_ids.get(req_id) + self.input_batch.remove_request(req_id) + # update the request state in self.input_batch + self.input_batch.add_request(req_state) + self._init_mrope_positions(req_state) + @torch.inference_mode() def execute_model( self, @@ -74,6 +85,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): record_function_or_nullcontext("gpu_model_runner: preprocess"), self.synchronize_input_prep(), ): + if self.model_config.async_chunk: + self._update_request_states(scheduler_output) self._update_states(scheduler_output) if not scheduler_output.total_num_scheduled_tokens: return EMPTY_MODEL_RUNNER_OUTPUT diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 06abd9f3a..a486212aa 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -808,6 +808,7 @@ class OmniGPUModelRunner(GPUModelRunner): hidden_states: torch.Tensor, multimodal_outputs: object, num_scheduled_tokens_np: np.ndarray, + scheduler_output: "SchedulerOutput", ) -> None: """Process model-provided per-request additional_information updates and merge into request state.""" try: @@ -815,10 +816,13 @@ class OmniGPUModelRunner(GPUModelRunner): # TODO(Peiqi): do we have a more elegant way to do this? if hasattr(self.model, "has_postprocess") and self.model.has_postprocess: for req_index, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests.get(req_id) - req_infos = ( - getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - ) + if self.model_config.async_chunk: + req_infos = self._get_additional_information(scheduler_output, req_id) + else: + req_state = self.requests.get(req_id) + req_infos = ( + getattr(req_state, "additional_information_cpu", None) if req_state is not None else None + ) start_offset = int(self.query_start_loc.cpu[req_index]) sched_tokens = int(num_scheduled_tokens_np[req_index]) s, e = start_offset, start_offset + sched_tokens @@ -859,6 +863,40 @@ class OmniGPUModelRunner(GPUModelRunner): start_offset = int(self.query_start_loc.cpu[req_index]) self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src) + def _get_additional_information(self, scheduler_output: "SchedulerOutput", req_id: str) -> dict: + req_infos = None + req_state = self.requests.get(req_id) + additional_information_cpu = getattr(req_state, "additional_information_cpu", None) + for new_req in scheduler_output.scheduled_new_reqs: + if new_req.req_id == req_id: + payload_info = getattr(new_req, "additional_information", None) + if payload_info is not None: + return payload_info + + if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"): + cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {}) + if isinstance(cached_infos, dict) and req_id in cached_infos: + req_infos = cached_infos[req_id] + if not isinstance(req_infos, dict): + req_infos = None + + if req_infos is None or req_infos.get("last_talker_hidden", None) is None: + if req_infos is None: + additional_information_cpu.pop("thinker_embeddings", None) + req_infos = additional_information_cpu + else: + req_infos["last_talker_hidden"] = additional_information_cpu.get("last_talker_hidden", None) + req_infos["num_processed_thinker_tokens"] = additional_information_cpu.get( + "num_processed_thinker_tokens", 0 + ) + if not isinstance(req_infos, dict): + req_infos = None + + if req_infos is None: + logger.warning(f"No additional_information found for req_id: {req_id}") + + return req_infos + def _preprocess( self, scheduler_output: "SchedulerOutput", @@ -974,9 +1012,14 @@ class OmniGPUModelRunner(GPUModelRunner): # collect additional_information (tensor/list) for prefill portion only decode_req_ids = [] for req_index, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests.get(req_id) - req_infos = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - + # Try to get additional_information from multiple sources + if self.vllm_config.model_config.async_chunk: + req_infos = self._get_additional_information(scheduler_output, req_id) + else: + req_state = self.requests.get(req_id) + req_infos = ( + getattr(req_state, "additional_information_cpu", None) if req_state is not None else None + ) start_offset = int(self.query_start_loc.cpu[req_index]) sched_tokens = int(num_scheduled_tokens_np[req_index]) s, e = start_offset, start_offset + sched_tokens