# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """OmniIntermediateBuffer — per-request cross-stage state for Omni pipelines. Uses ``req_index`` (not ``req_id``) for O(1) access, aligned with v2's ``RequestState`` slot management. """ from __future__ import annotations from typing import Any import numpy as np import torch from vllm.v1.core.sched.output import NewRequestData from vllm.v1.worker.gpu.input_batch import InputBatch from vllm_omni.worker.payload_span import get_tensor_span, merge_tensor_spans def _resolve_prompt_embeds(pe: Any) -> torch.Tensor | None: """Convert a prompt_embeds payload to a contiguous CPU tensor.""" if pe is None: return None if isinstance(pe, torch.Tensor): return pe.detach().cpu().contiguous() try: data = getattr(pe, "data", None) shape = getattr(pe, "shape", None) if data is not None and shape is not None: dt = np.dtype(getattr(pe, "dtype", "float32")) arr = np.frombuffer(data, dtype=dt).reshape(shape) return torch.from_numpy(arr.copy()) except Exception as exc: raise ValueError("Failed to decode prompt_embeds payload") from exc raise TypeError(f"Unsupported prompt_embeds payload type: {type(pe).__name__}") def _resolve_additional_information(payload: Any) -> dict[str, Any]: """Convert an additional_information payload to a plain dict.""" if payload is None: return {} try: from vllm_omni.data_entry_keys import unflatten_payload from vllm_omni.engine.serialization import deserialize_additional_information info = deserialize_additional_information(payload) if isinstance(info, dict) and any(isinstance(key, str) and "." in key for key in info): info = unflatten_payload(info) except Exception as exc: raise ValueError("Failed to decode additional_information payload") from exc if not isinstance(info, dict): raise TypeError(f"Decoded additional_information must be a dict, got {type(info).__name__}") return info class OmniIntermediateBuffer: """Per-request intermediate state for multi-stage Omni pipelines. Stores ``prompt_embeds``, ``additional_information``, ``mm_features``, ``req_id`` and any runtime updates written back by model postprocess. """ def __init__(self, max_num_reqs: int): self.buffers: list[dict[str, Any]] = [{} for _ in range(max_num_reqs)] self.req_id_to_index: dict[str, int] = {} def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: info: dict[str, Any] = {} pe = getattr(new_req_data, "prompt_embeds", None) if pe is not None: pe_cpu = _resolve_prompt_embeds(pe) if pe_cpu is not None: info["prompt_embeds_cpu"] = pe_cpu ai = getattr(new_req_data, "additional_information", None) if ai is not None: info.update(_resolve_additional_information(ai)) if new_req_data.mm_features: info["mm_features"] = new_req_data.mm_features sampling_params = getattr(new_req_data, "sampling_params", None) if sampling_params is not None: info["sampling_params"] = sampling_params info["req_id"] = new_req_data.req_id previous_id = self.buffers[req_index].get("req_id") if previous_id is not None and self.req_id_to_index.get(previous_id) == req_index: self.req_id_to_index.pop(previous_id, None) self.buffers[req_index] = info self.req_id_to_index[new_req_data.req_id] = req_index def remove_request(self, req_index: int) -> None: req_id = self.buffers[req_index].get("req_id") if req_id is not None and self.req_id_to_index.get(req_id) == req_index: self.req_id_to_index.pop(req_id) self.buffers[req_index] = {} def gather(self, input_batch: InputBatch) -> list[dict[str, Any]]: """Return buffer dicts in current batch order (via ``idx_mapping_np``).""" return [self.buffers[idx] for idx in input_batch.idx_mapping_np] @staticmethod def _split_gpu_keys(gpu_resident_keys: set[Any] | None) -> tuple[set[Any], dict[Any, set[Any]]]: top_level: set[Any] = set() nested: dict[Any, set[Any]] = {} for key in gpu_resident_keys or set(): if isinstance(key, tuple) and len(key) == 2: nested.setdefault(key[0], set()).add(key[1]) else: top_level.add(key) return top_level, nested @staticmethod def _store_value(dest: dict[Any, Any], key: Any, value: Any, gpu_keys: set[Any]) -> None: if isinstance(value, torch.Tensor): if key in gpu_keys: dest[key] = value.detach().clone() else: dest[key] = value.detach().cpu().contiguous() elif isinstance(value, list): dest[key] = [ ( item.detach().clone() if key in gpu_keys and isinstance(item, torch.Tensor) else item.detach().cpu().contiguous() if isinstance(item, torch.Tensor) else item ) for item in value ] else: dest[key] = value def update( self, req_index: int, updates: dict[Any, Any], gpu_resident_keys: set[Any] | None = None, ) -> None: """Merge *updates* into the buffer at *req_index*. Tensors are detached; those whose key is **not** in *gpu_resident_keys* are moved to CPU. """ if not updates: return top_gpu_keys, nested_gpu_keys = self._split_gpu_keys(gpu_resident_keys) existing = self.buffers[req_index] for k, v in updates.items(): if isinstance(k, tuple) and len(k) == 2: type_key, qualifier = k existing_sub = existing.setdefault(type_key, {}) if not isinstance(existing_sub, dict): existing_sub = {} existing[type_key] = existing_sub self._store_value(existing_sub, qualifier, v, nested_gpu_keys.get(type_key, set())) elif isinstance(v, dict): existing_sub = existing.setdefault(k, {}) if not isinstance(existing_sub, dict): existing_sub = {} existing[k] = existing_sub merged_qualifiers: set[str] = set() if k == "embed" and isinstance(v.get("decode"), torch.Tensor): span_keys = {"decode_token_start", "decode_token_end"} if span_keys.intersection(v): incoming_span = get_tensor_span( v, tensor_key="decode", start_key="decode_token_start", end_key="decode_token_end", ) if incoming_span is None: raise RuntimeError("Incoming Thinker decode delta has an invalid absolute span") existing_span = get_tensor_span( existing_sub, tensor_key="decode", start_key="decode_token_start", end_key="decode_token_end", ) merged_span: tuple[Any, int, int] | None = incoming_span if existing_span is not None: merged_span = merge_tensor_spans(existing_span, incoming_span) if merged_span is None: raise RuntimeError( "Non-contiguous pending Thinker decode spans in model buffer: " f"existing=[{existing_span[1]}, {existing_span[2]}) " f"incoming=[{incoming_span[1]}, {incoming_span[2]})" ) assert merged_span is not None merged_tensor, merged_start, merged_end = merged_span self._store_value( existing_sub, "decode", merged_tensor, nested_gpu_keys.get(k, set()), ) existing_sub["decode_token_start"] = merged_start existing_sub["decode_token_end"] = merged_end merged_qualifiers.update({"decode", "decode_token_start", "decode_token_end"}) for qualifier, value in v.items(): if qualifier in merged_qualifiers: continue self._store_value(existing_sub, qualifier, value, nested_gpu_keys.get(k, set())) else: self._store_value(existing, k, v, top_gpu_keys) def update_gpu_tensor_rows( self, req_indices: list[int], key: Any, values: torch.Tensor, *, keepdim: bool = True, ) -> None: """Snapshot a batch-first tensor once and retain owned row views.""" self._store_gpu_tensor_rows(req_indices, key, values.detach().clone(), keepdim=keepdim) def update_owned_gpu_tensor_rows( self, req_indices: list[int], key: Any, owned_values: torch.Tensor, *, keepdim: bool = True, ) -> None: """Store row views of an already-owned batch tensor without cloning. The caller must guarantee ``owned_values`` is not written after this call (the model output ownership contract). Row views keep the owned storage alive until each request row is replaced or the slot is freed. """ self._store_gpu_tensor_rows(req_indices, key, owned_values, keepdim=keepdim) def _store_gpu_tensor_rows( self, req_indices: list[int], key: Any, owned: torch.Tensor, *, keepdim: bool, ) -> None: num_rows = len(req_indices) if owned.ndim == 0 or owned.shape[0] != num_rows: raise ValueError( "Batched GPU state update changed the request axis: " f"key={key!r} expected={num_rows} " f"actual={owned.shape[0] if owned.ndim else 0}" ) for row, req_index in enumerate(req_indices): existing = self.buffers[req_index] row_value = owned[row : row + 1] if keepdim else owned[row] if isinstance(key, tuple) and len(key) == 2: type_key, qualifier = key existing_sub = existing.setdefault(type_key, {}) if not isinstance(existing_sub, dict): existing_sub = {} existing[type_key] = existing_sub existing_sub[qualifier] = row_value else: existing[key] = row_value