282 lines
12 KiB
Python
282 lines
12 KiB
Python
# 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))
|
|
|
|
# Orchestrator stage bridges (e.g. MiniCPM-o's llm2tts) hand the
|
|
# request's runner payload over directly, like the V1 runner's
|
|
# ``_update_intermediate_buffer``: nested sections merge key by key.
|
|
mib = getattr(new_req_data, "model_intermediate_buffer", None)
|
|
if isinstance(mib, dict) and mib:
|
|
for key, value in mib.items():
|
|
if isinstance(value, dict):
|
|
section = info.get(key)
|
|
if not isinstance(section, dict):
|
|
section = {}
|
|
info[key] = section
|
|
for qualifier, item in value.items():
|
|
self._store_value(section, qualifier, item, set())
|
|
else:
|
|
self._store_value(info, key, value, set())
|
|
|
|
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
|