[cpu worker] file move

Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>

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)

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__)
 

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()
+        gc.collect()

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__)

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__)
 

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):

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__)
 

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):