From 8ee1d73d025702267e13e6a6a67a7eaa49b22d2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=BD=90=E4=BF=9D=E5=85=83?= Date: Fri, 14 Nov 2025 17:38:55 +0800 Subject: [PATCH] [cpu worker] file move MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 齐保元 --- vllm_omni/worker/cpu/__init__.py | 0 .../worker/{ => cpu}/cpu_ar_model_runner.py | 171 ++++++++++++------ vllm_omni/worker/{ => cpu}/cpu_ar_worker.py | 2 +- .../{ => cpu}/cpu_diffusion_model_runner.py | 10 +- .../worker/{ => cpu}/cpu_diffusion_worker.py | 2 +- .../worker/{ => cpu}/cpu_model_runner.py | 0 vllm_omni/worker/gpu/__init__.py | 0 .../worker/{ => gpu}/gpu_ar_model_runner.py | 2 +- vllm_omni/worker/{ => gpu}/gpu_ar_worker.py | 2 +- .../{ => gpu}/gpu_diffusion_model_runner.py | 2 +- .../worker/{ => gpu}/gpu_diffusion_worker.py | 2 +- .../worker/{ => gpu}/gpu_model_runner.py | 0 12 files changed, 127 insertions(+), 66 deletions(-) create mode 100644 vllm_omni/worker/cpu/__init__.py rename vllm_omni/worker/{ => cpu}/cpu_ar_model_runner.py (77%) rename vllm_omni/worker/{ => cpu}/cpu_ar_worker.py (97%) rename vllm_omni/worker/{ => cpu}/cpu_diffusion_model_runner.py (99%) rename vllm_omni/worker/{ => cpu}/cpu_diffusion_worker.py (96%) rename vllm_omni/worker/{ => cpu}/cpu_model_runner.py (100%) create mode 100644 vllm_omni/worker/gpu/__init__.py rename vllm_omni/worker/{ => gpu}/gpu_ar_model_runner.py (99%) rename vllm_omni/worker/{ => gpu}/gpu_ar_worker.py (97%) rename vllm_omni/worker/{ => gpu}/gpu_diffusion_model_runner.py (99%) rename vllm_omni/worker/{ => gpu}/gpu_diffusion_worker.py (97%) rename vllm_omni/worker/{ => gpu}/gpu_model_runner.py (100%) diff --git a/vllm_omni/worker/cpu/__init__.py b/vllm_omni/worker/cpu/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/worker/cpu_ar_model_runner.py b/vllm_omni/worker/cpu/cpu_ar_model_runner.py similarity index 77% rename from vllm_omni/worker/cpu_ar_model_runner.py rename to vllm_omni/worker/cpu/cpu_ar_model_runner.py index 5ee44cd2d..f4ce7f5dd 100644 --- a/vllm_omni/worker/cpu_ar_model_runner.py +++ b/vllm_omni/worker/cpu/cpu_ar_model_runner.py @@ -1,4 +1,4 @@ -"""AR CPU Model Runner for vLLM-omni. +"""AR GPU Model Runner for vLLM-omni. Exposes per-request hidden representations via ModelRunnerOutput.pooler_output and also outputs sampled tokens. @@ -30,13 +30,20 @@ from vllm.v1.worker.ubatch_utils import UBatchSlices from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload from vllm_omni.outputs import OmniModelRunnerOutput -from vllm_omni.worker.cpu_model_runner import OmniCPUModelRunner +from vllm_omni.worker.cpu.cpu_model_runner import OmniCPUModelRunner logger = init_logger(__name__) class CPUARModelRunner(OmniCPUModelRunner): + """Autoregressive GPU model runner that returns hidden states per request. + This runner follows the same preparation and forward path as GPUModelRunner + (inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs), + and additionally performs lightweight sampling so that sampled tokens are + available in outputs. Hidden representations are taken at the same indices + that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`). + """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -58,6 +65,7 @@ class CPUARModelRunner(OmniCPUModelRunner): torch.Tensor, Optional[IntermediateTensors], dict[str, Any], + Optional[dict[str, dict]], ]: num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens @@ -72,6 +80,7 @@ class CPUARModelRunner(OmniCPUModelRunner): # _prepare_inputs may reorder the batch, so we must gather multi # modal outputs after that to ensure the correct order + per_req_additional_information: Optional[dict[str, dict]] = None if ( self.supports_mm_inputs and get_pp_group().is_first_rank @@ -92,9 +101,11 @@ class CPUARModelRunner(OmniCPUModelRunner): # TODO(woosuk): Avoid the copy. Optimize. self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) - # Reset per-step additional information collector + # Reset per-step additional information collector (deprecated concat path) if hasattr(self, "_forward_additional_information"): self._forward_additional_information = None + # New: per-request additional information for this step + per_req_additional_information = {} # Overlay custom prompt_embeds per request for the prompt portion; # collect additional_information (tensor/list) for prefill portion only @@ -115,54 +126,31 @@ class CPUARModelRunner(OmniCPUModelRunner): src = pe_cpu[ num_computed_tokens : num_computed_tokens + overlay_len ].to(dtype=self.dtype, device=self.device, non_blocking=True) - start_offset = int(self.query_start_loc_cpu[req_index]) + start_offset = int(self.query_start_loc.cpu[req_index]) self.inputs_embeds[start_offset : start_offset + overlay_len].copy_( src ) - # For additional_information: handle arbitrary keys + # Build per-request additional information (no cross-request concat) if addi_cpu is not None and isinstance(addi_cpu, dict): - # Lazy init collector dict - if ( - not hasattr(self, "_forward_additional_information") - or self._forward_additional_information is None - ): - self._forward_additional_information = {} - # Process tensors (slice by scheduled prompt range) and - # lists (append per-request) + req_info: dict[str, object] = {} for k, v in addi_cpu.items(): if isinstance(v, torch.Tensor): - # Slice along token dimension for prefill part - try: - seg = v[ - num_computed_tokens : num_computed_tokens - + overlay_len - ].to( - dtype=self.dtype, - device=self.device, - non_blocking=True, - ) - except Exception: - # Fallback: move whole tensor if slicing fails - seg = v.to( - dtype=self.dtype, - device=self.device, - non_blocking=True, - ) - prev_val = self._forward_additional_information.get(k) - self._forward_additional_information[k] = ( - torch.cat([prev_val, seg], dim=0) - if isinstance(prev_val, torch.Tensor) - else seg.clone() - ) - elif isinstance(v, list): - prev_val = self._forward_additional_information.get(k) - if prev_val is None: - self._forward_additional_information[k] = [v] - elif isinstance(prev_val, list): - self._forward_additional_information[k].append(v) + # For prefill tokens, pass only the scheduled slice; for decode or no scheduled tokens, pass whole tensor + if overlay_len > 0: + try: + seg = v[ + num_computed_tokens : num_computed_tokens + overlay_len + ].detach().to("cpu").contiguous() + except Exception: + seg = v.detach().to("cpu").contiguous() + req_info[k] = seg else: - # Mixed types: wrap existing into list - self._forward_additional_information[k] = [prev_val, v] + req_info[k] = v.detach().to("cpu").contiguous() + elif isinstance(v, list): + req_info[k] = v + else: + req_info[k] = v + per_req_additional_information[req_id] = req_info input_ids = self.input_ids.gpu[:num_input_tokens] inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] @@ -233,6 +221,7 @@ class CPUARModelRunner(OmniCPUModelRunner): positions, intermediate_tensors, model_kwargs, + per_req_additional_information, ) @torch.inference_mode() @@ -348,6 +337,7 @@ class CPUARModelRunner(OmniCPUModelRunner): positions, intermediate_tensors, model_kwargs, + per_req_additional_information, ) = self._preprocess( scheduler_output, num_scheduled_tokens_np, @@ -388,15 +378,33 @@ class CPUARModelRunner(OmniCPUModelRunner): ): model_kwargs_extra = {} - # Only pass additional_information for the prefill part - if ( - hasattr(self, "_forward_additional_information") - and self._forward_additional_information is not None - and isinstance(self._forward_additional_information, dict) - ): - model_kwargs_extra["additional_information"] = ( - self._forward_additional_information + # Pass per-request additional information map for this step (no concat) + if per_req_additional_information: + model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information + # Always pass per-request runtime additional_information (persisted in request state) + try: + per_req_runtime_info = [] + for req_id in self.input_batch.req_ids: + req_state = self.requests.get(req_id) + info = ( + getattr(req_state, "additional_information_cpu", None) + if req_state is not None + else None + ) + per_req_runtime_info.append(info if isinstance(info, dict) else {}) + model_kwargs_extra["runtime_additional_information"] = ( + per_req_runtime_info ) + model_kwargs_extra["request_ids"] = self.input_batch.req_ids + # Pass each request's token span within the flattened sequence for this step, enabling the model to map decode/prefill by request + req_token_spans = [] + for req_index in range(len(self.input_batch.req_ids)): + start_offset = int(self.query_start_loc.cpu[req_index]) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + req_token_spans.append((start_offset, start_offset + sched_tokens)) + model_kwargs_extra["request_token_spans"] = req_token_spans + except Exception: + pass model_output = self.model( input_ids=input_ids, positions=positions, @@ -421,8 +429,42 @@ class CPUARModelRunner(OmniCPUModelRunner): hidden_states, multimodal_outputs = self.extract_multimodal_outputs( hidden_states ) - + # The model side may return per-request additional_information updates (model-agnostic channel). + # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order; + # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode. + try: + if ( + isinstance(multimodal_outputs, dict) + and ( + "additional_information_update" in multimodal_outputs + or "additional_information_update_by_req_id" in multimodal_outputs + ) + ): + # Option A: list[dict] in batch order + updates_list = multimodal_outputs.get( + "additional_information_update" + ) + if isinstance(updates_list, list): + for idx, upd in enumerate(updates_list): + if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids): + continue + req_id = self.input_batch.req_ids[idx] + self._merge_additional_information_update(req_id, upd) + # Option B: dict[str, dict] keyed by req_id + updates_map = multimodal_outputs.get( + "additional_information_update_by_req_id" + ) + if isinstance(updates_map, dict): + for req_id, upd in updates_map.items(): + if not isinstance(upd, dict): + continue + if req_id not in self.requests: + continue + self._merge_additional_information_update(req_id, upd) + except Exception as e: + logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}, with the multimodal_outputs as {multimodal_outputs}") if not self.broadcast_pp_output: + # Common case. if not get_pp_group().is_last_rank: # Return the intermediate tensors. @@ -582,3 +624,26 @@ class CPUARModelRunner(OmniCPUModelRunner): invalid_req_indices=invalid_req_indices, async_output_copy_stream=self.async_output_copy_stream, ) + + def _merge_additional_information_update(self, req_id: str, upd: dict) -> None: + req_state = self.requests.get(req_id) + if req_state is None: + return + existing = getattr(req_state, "additional_information_cpu", {}) + if not isinstance(existing, dict): + existing = {} + merged = dict(existing) + for k, v in upd.items(): + if isinstance(v, torch.Tensor): + merged[k] = v.detach().to("cpu").contiguous() + elif isinstance(v, list): + new_list = [] + for item in v: + if isinstance(item, torch.Tensor): + new_list.append(item.detach().to("cpu").contiguous()) + else: + new_list.append(item) + merged[k] = new_list + else: + merged[k] = v + setattr(req_state, "additional_information_cpu", merged) \ No newline at end of file diff --git a/vllm_omni/worker/cpu_ar_worker.py b/vllm_omni/worker/cpu/cpu_ar_worker.py similarity index 97% rename from vllm_omni/worker/cpu_ar_worker.py rename to vllm_omni/worker/cpu/cpu_ar_worker.py index 7a13de5b8..50b42e400 100644 --- a/vllm_omni/worker/cpu_ar_worker.py +++ b/vllm_omni/worker/cpu/cpu_ar_worker.py @@ -11,7 +11,7 @@ from vllm.platforms import current_platform, CpuArchEnum from vllm.v1.worker.cpu_worker import CPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment -from vllm_omni.worker.cpu_ar_model_runner import CPUARModelRunner +from vllm_omni.worker.cpu.cpu_ar_model_runner import CPUARModelRunner logger = init_logger(__name__) diff --git a/vllm_omni/worker/cpu_diffusion_model_runner.py b/vllm_omni/worker/cpu/cpu_diffusion_model_runner.py similarity index 99% rename from vllm_omni/worker/cpu_diffusion_model_runner.py rename to vllm_omni/worker/cpu/cpu_diffusion_model_runner.py index 43b684e31..ba4edcbf3 100644 --- a/vllm_omni/worker/cpu_diffusion_model_runner.py +++ b/vllm_omni/worker/cpu/cpu_diffusion_model_runner.py @@ -10,11 +10,7 @@ import torch from vllm.config import CUDAGraphMode from vllm.forward_context import BatchDescriptor from vllm.multimodal.inputs import MultiModalKwargs -from vllm.utils import cdiv -from vllm.v1.attention.backends.utils import ( - CommonAttentionMetadata, - split_attn_metadata, -) +from vllm.v1.attention.backends.utils import CommonAttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.utils import record_function_or_nullcontext @@ -29,7 +25,7 @@ from vllm.v1.worker.gpu_model_runner import ( from vllm.v1.worker.ubatch_utils import UBatchSlices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs from vllm_omni.outputs import OmniModelRunnerOutput -from vllm_omni.worker.cpu_model_runner import OmniCPUModelRunner +from vllm_omni.worker.cpu.cpu_model_runner import OmniCPUModelRunner logger = logging.getLogger(__name__) @@ -621,4 +617,4 @@ class CPUDiffusionModelRunner(OmniCPUModelRunner): self._sync_device() del hidden_states self.encoder_cache.clear() - gc.collect() \ No newline at end of file + gc.collect() diff --git a/vllm_omni/worker/cpu_diffusion_worker.py b/vllm_omni/worker/cpu/cpu_diffusion_worker.py similarity index 96% rename from vllm_omni/worker/cpu_diffusion_worker.py rename to vllm_omni/worker/cpu/cpu_diffusion_worker.py index b4db58ff7..c94c73c0c 100644 --- a/vllm_omni/worker/cpu_diffusion_worker.py +++ b/vllm_omni/worker/cpu/cpu_diffusion_worker.py @@ -13,7 +13,7 @@ from vllm.v1.worker.cpu_worker import CPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment -from vllm_omni.worker.cpu_diffusion_model_runner import CPUDiffusionModelRunner +from vllm_omni.worker.cpu.cpu_diffusion_model_runner import CPUDiffusionModelRunner logger = init_logger(__name__) diff --git a/vllm_omni/worker/cpu_model_runner.py b/vllm_omni/worker/cpu/cpu_model_runner.py similarity index 100% rename from vllm_omni/worker/cpu_model_runner.py rename to vllm_omni/worker/cpu/cpu_model_runner.py diff --git a/vllm_omni/worker/gpu/__init__.py b/vllm_omni/worker/gpu/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu/gpu_ar_model_runner.py similarity index 99% rename from vllm_omni/worker/gpu_ar_model_runner.py rename to vllm_omni/worker/gpu/gpu_ar_model_runner.py index d91defd81..99b9beb30 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu/gpu_ar_model_runner.py @@ -30,7 +30,7 @@ from vllm.v1.worker.ubatch_utils import UBatchSlices from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload from vllm_omni.outputs import OmniModelRunnerOutput -from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner +from vllm_omni.worker.gpu.gpu_model_runner import OmniGPUModelRunner logger = init_logger(__name__) diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu/gpu_ar_worker.py similarity index 97% rename from vllm_omni/worker/gpu_ar_worker.py rename to vllm_omni/worker/gpu/gpu_ar_worker.py index 838068df7..c4b30132b 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu/gpu_ar_worker.py @@ -9,7 +9,7 @@ from vllm.utils import GiB_bytes, MemorySnapshot from vllm.v1.utils import report_usage_stats from vllm.v1.worker.gpu_worker import Worker as GPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment -from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner +from vllm_omni.worker.gpu.gpu_ar_model_runner import GPUARModelRunner class GPUARWorker(GPUWorker): diff --git a/vllm_omni/worker/gpu_diffusion_model_runner.py b/vllm_omni/worker/gpu/gpu_diffusion_model_runner.py similarity index 99% rename from vllm_omni/worker/gpu_diffusion_model_runner.py rename to vllm_omni/worker/gpu/gpu_diffusion_model_runner.py index e8cd68491..0e75e4fc9 100644 --- a/vllm_omni/worker/gpu_diffusion_model_runner.py +++ b/vllm_omni/worker/gpu/gpu_diffusion_model_runner.py @@ -25,7 +25,7 @@ from vllm.v1.worker.gpu_model_runner import ( from vllm.v1.worker.ubatch_utils import UBatchSlices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs from vllm_omni.outputs import OmniModelRunnerOutput -from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner +from vllm_omni.worker.gpu.gpu_model_runner import OmniGPUModelRunner logger = logging.getLogger(__name__) diff --git a/vllm_omni/worker/gpu_diffusion_worker.py b/vllm_omni/worker/gpu/gpu_diffusion_worker.py similarity index 97% rename from vllm_omni/worker/gpu_diffusion_worker.py rename to vllm_omni/worker/gpu/gpu_diffusion_worker.py index 46af50190..4dae8292a 100644 --- a/vllm_omni/worker/gpu_diffusion_worker.py +++ b/vllm_omni/worker/gpu/gpu_diffusion_worker.py @@ -9,7 +9,7 @@ from vllm.utils import GiB_bytes, MemorySnapshot from vllm.v1.utils import report_usage_stats from vllm.v1.worker.gpu_worker import Worker as GPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment -from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner +from vllm_omni.worker.gpu.gpu_diffusion_model_runner import GPUDiffusionModelRunner class GPUDiffusionWorker(GPUWorker): diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu/gpu_model_runner.py similarity index 100% rename from vllm_omni/worker/gpu_model_runner.py rename to vllm_omni/worker/gpu/gpu_model_runner.py