Yueqian LinandRein Yang
e996723bc9
Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Signed-off-by: Rein Yang <ruiruyang2@gmail.com> Co-authored-by: Rein Yang <ruiruyang2@gmail.com> the ci failure case is irrelavent to this PR and occurs rarely
452 lines
19 KiB
Python
452 lines
19 KiB
Python
"""
|
|
Utilities for Prefix Caching in Omni models.
|
|
"""
|
|
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
from vllm.v1.worker.gpu_input_batch import InputBatch
|
|
|
|
from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class OmniTensorPrefixCache:
|
|
"""Prefix cache for hidden states (model outputs) and model specific
|
|
multimodal outputs.
|
|
|
|
This class implements prefix caching in a non-invasive way on top of
|
|
vLLM by leveraging the same slot mappings that the vLLM scheduler uses
|
|
for the KV Cache.
|
|
|
|
Conceptually, this means we are mapping vLLM's cache mapping:
|
|
(num_blocks, block_size)
|
|
|
|
to 3D tensors of shape:
|
|
(num_blocks, block_size, feature_size)
|
|
|
|
Note that feature_size may vary across multimodal_outputs.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
num_blocks: int,
|
|
block_size: int,
|
|
hidden_size: int,
|
|
hs_dtype: torch.dtype,
|
|
):
|
|
self.num_blocks = num_blocks
|
|
self.block_size = block_size
|
|
self.default_hidden_size = hidden_size
|
|
|
|
# Initialize the hidden states cache immediately
|
|
self.hidden_states_cache = self._get_cache_tensor(dtype=hs_dtype)
|
|
|
|
# Defer initialization of the mm_outputs_cache until we
|
|
# actually see mm output tensors dependent on num tokens.
|
|
self.mm_outputs_cache = {}
|
|
self.mm_cache_keys = set()
|
|
self._new_req_cache_hit_ids: set[str] = set()
|
|
self._deferred_mm_outputs: dict[str, dict[str, list[tuple[int, torch.Tensor]]]] = {}
|
|
|
|
def maybe_init_missing_mm_cache_keys(self, multimodal_outputs: dict, seq_len: int):
|
|
"""Given multimodal outputs from executing the model, dynamically
|
|
determine which multimodal outputs are tensors depending on sequence
|
|
length and should be cached, and initialize the cache tensors
|
|
accordingly.
|
|
|
|
NOTE: This is done to avoid the need for explicit specification of
|
|
cache keys for every model/stage and aligns with the current way
|
|
that we slice the multimodal outputs based on the first dimension.
|
|
|
|
This will usually be called by the first forward pass, i.e.,
|
|
determined by the warmup.
|
|
"""
|
|
for key, val in multimodal_outputs.items():
|
|
# Only cache per-token feature tensors: 2D+ with first dim == seq_len.
|
|
# A 1D tensor of shape (seq_len,) is a broadcast scalar (per-request
|
|
# metadata such as ref_code_len / codec_streaming), not per-token data;
|
|
# caching it by slot causes a shape mismatch when a later request has a
|
|
# different scheduled seq length.
|
|
if (
|
|
isinstance(val, torch.Tensor)
|
|
and val.ndim >= 2
|
|
and val.shape[0] == seq_len
|
|
and key not in self.mm_cache_keys
|
|
):
|
|
feat_dim = val.shape[-1]
|
|
self.mm_outputs_cache[key] = self._get_cache_tensor(
|
|
dtype=val.dtype,
|
|
hidden_size=feat_dim,
|
|
)
|
|
self.mm_cache_keys.add(key)
|
|
new_tensor_shape = self.mm_outputs_cache[key].shape
|
|
logger.info("Initializing multimodal output cache of size %s for key: %s", list(new_tensor_shape), key)
|
|
|
|
def _get_cache_tensor(self, dtype: torch.dtype, hidden_size: int | None = None) -> torch.Tensor:
|
|
"""Allocate a CPU cache tensor for a specific key."""
|
|
actual_hidden_size = hidden_size if hidden_size is not None else self.default_hidden_size
|
|
return torch.zeros(
|
|
(self.num_blocks, self.block_size, actual_hidden_size),
|
|
dtype=dtype,
|
|
device="cpu",
|
|
)
|
|
|
|
def add_prefix_cached_new_req_id(self, req_id: str):
|
|
"""Adds a new request ID to the set of prefix cache hits on the batch."""
|
|
self._new_req_cache_hit_ids.add(req_id)
|
|
|
|
def reset_prefix_cached_new_req_ids(self):
|
|
"""Clears the cache hit IDs to prepare for a new engine step."""
|
|
self._new_req_cache_hit_ids.clear()
|
|
|
|
def has_prefix_cached_new_req_ids(self) -> bool:
|
|
"""Return True when this step contains a newly scheduled prefix hit."""
|
|
return bool(self._new_req_cache_hit_ids)
|
|
|
|
@staticmethod
|
|
def _coerce_to_cpu_tensor(maybe_gpu_tensor: torch.Tensor) -> torch.Tensor:
|
|
"""Convert GPU tensors -> contiguous CPU tensors if needed."""
|
|
return maybe_gpu_tensor.detach().cpu().contiguous()
|
|
|
|
def update_omni_tensor_prefix_cache(
|
|
self,
|
|
hidden_states: torch.Tensor | None,
|
|
multimodal_outputs: dict[str, torch.Tensor] | None,
|
|
num_tokens_unpadded: int,
|
|
slot_mapping: torch.Tensor,
|
|
num_tokens_padded: int | None = None,
|
|
skip_mm_cache_keys: set[str] | None = None,
|
|
):
|
|
"""Updates the hidden cache state for the provided hidden states and multimodal outputs.
|
|
|
|
Args:
|
|
hidden_states: Hidden states tensor to cache (if any)
|
|
multimodal_outputs: Multimodal dict whose tensors may be cached
|
|
num_tokens_unpadded: Number of tokens without padding
|
|
slot_mapping: Slot mapping for the input sequence
|
|
num_tokens_padded: Total number of tokens including padding
|
|
skip_mm_cache_keys: Multimodal keys whose CPU cache writes are deferred
|
|
"""
|
|
unpadded_slot_mapping = slot_mapping[:num_tokens_unpadded]
|
|
if num_tokens_padded is None:
|
|
num_tokens_padded = num_tokens_unpadded
|
|
skip_mm_cache_keys = skip_mm_cache_keys or set()
|
|
|
|
if hidden_states is not None:
|
|
# Slice to unpadded portion before caching
|
|
hidden_states = hidden_states[:num_tokens_unpadded]
|
|
# Ensure that hidden states are on the CPU
|
|
hidden_states = OmniTensorPrefixCache._coerce_to_cpu_tensor(hidden_states)
|
|
# View the cache as 2D so that we can treat our slots as row indices
|
|
flat_cache = self.hidden_states_cache.view(-1, self.hidden_states_cache.shape[-1])
|
|
flat_cache[unpadded_slot_mapping] = hidden_states
|
|
logger.debug("Writing to hidden states for %s tokens", num_tokens_unpadded)
|
|
|
|
# Do the same for the stage's cached multimodal outputs
|
|
if multimodal_outputs is not None:
|
|
# If we haven't initialized the keys already, do it now
|
|
# We check against the padded token count since we haven't sliced yet
|
|
self.maybe_init_missing_mm_cache_keys(
|
|
multimodal_outputs,
|
|
seq_len=num_tokens_padded,
|
|
)
|
|
|
|
for mm_out_key, mm_cache in self.mm_outputs_cache.items():
|
|
if mm_out_key in multimodal_outputs:
|
|
if mm_out_key in skip_mm_cache_keys:
|
|
continue
|
|
# Slice to unpadded portion before caching
|
|
mm_state = multimodal_outputs[mm_out_key][:num_tokens_unpadded]
|
|
mm_state = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state)
|
|
flat_cache = mm_cache.view(-1, mm_cache.shape[-1])
|
|
flat_cache[unpadded_slot_mapping] = mm_state
|
|
logger.debug("Writing to mm output cache for %s tokens", num_tokens_unpadded)
|
|
|
|
def stage_deferred_mm_outputs(
|
|
self,
|
|
query_start_loc: torch.Tensor,
|
|
input_batch: InputBatch,
|
|
multimodal_outputs: dict[str, torch.Tensor] | None,
|
|
num_scheduled_tokens: dict[str, int],
|
|
deferred_mm_cache_keys: set[str],
|
|
) -> None:
|
|
"""Keep GPU multimodal chunks until a request finishes.
|
|
|
|
The normal prefix-cache write path copies every cached multimodal tensor
|
|
to CPU on every step. For model outputs that are only needed by future
|
|
full-block prefix hits, we can keep detached GPU chunks and materialize
|
|
the CPU cache once when the request completes.
|
|
"""
|
|
if not multimodal_outputs or not deferred_mm_cache_keys:
|
|
return
|
|
|
|
for mm_key in deferred_mm_cache_keys:
|
|
mm_state = multimodal_outputs.get(mm_key)
|
|
if not isinstance(mm_state, torch.Tensor) or mm_state.ndim < 2:
|
|
continue
|
|
for req_id in input_batch.req_ids:
|
|
req_idx = input_batch.req_id_to_index[req_id]
|
|
sched = int(num_scheduled_tokens.get(req_id, 0))
|
|
if sched <= 0:
|
|
continue
|
|
start = int(query_start_loc[req_idx])
|
|
end = start + sched
|
|
if start >= int(mm_state.shape[0]):
|
|
continue
|
|
end = min(end, int(mm_state.shape[0]))
|
|
if end <= start:
|
|
continue
|
|
|
|
computed_end = int(input_batch.num_computed_tokens_cpu[req_idx])
|
|
token_start = max(0, computed_end - sched)
|
|
chunk = mm_state[start:end].detach()
|
|
self._deferred_mm_outputs.setdefault(req_id, {}).setdefault(mm_key, []).append((token_start, chunk))
|
|
|
|
def commit_deferred_mm_outputs(
|
|
self,
|
|
finished_req_ids: set[str] | list[str],
|
|
input_batch: InputBatch,
|
|
) -> None:
|
|
"""Write deferred multimodal chunks into the CPU prefix cache.
|
|
|
|
This must run before finished requests are removed from ``input_batch``,
|
|
because the block table is needed to map logical token positions to KV
|
|
cache slots.
|
|
"""
|
|
if not finished_req_ids or not self._deferred_mm_outputs:
|
|
return
|
|
|
|
for req_id in tuple(finished_req_ids):
|
|
per_key = self._deferred_mm_outputs.pop(req_id, None)
|
|
if not per_key:
|
|
continue
|
|
req_idx = input_batch.req_id_to_index.get(req_id)
|
|
if req_idx is None:
|
|
continue
|
|
for mm_key, chunks in per_key.items():
|
|
self._commit_deferred_mm_key(req_idx, input_batch, mm_key, chunks)
|
|
|
|
def discard_deferred_mm_outputs(self, req_id: str) -> None:
|
|
"""Drop deferred chunks for requests that leave without a cache commit."""
|
|
self._deferred_mm_outputs.pop(req_id, None)
|
|
|
|
def _commit_deferred_mm_key(
|
|
self,
|
|
req_idx: int,
|
|
input_batch: InputBatch,
|
|
mm_key: str,
|
|
chunks: list[tuple[int, torch.Tensor]],
|
|
) -> None:
|
|
if not chunks:
|
|
return
|
|
|
|
chunks = sorted(chunks, key=lambda item: item[0])
|
|
first_start = chunks[0][0]
|
|
expected_start = first_start
|
|
tensors: list[torch.Tensor] = []
|
|
for token_start, tensor in chunks:
|
|
if tensor.numel() == 0:
|
|
continue
|
|
if token_start != expected_start:
|
|
if tensors:
|
|
self._write_deferred_mm_tensor(req_idx, input_batch, mm_key, tensors, first_start)
|
|
first_start = token_start
|
|
expected_start = token_start
|
|
tensors = []
|
|
tensors.append(tensor)
|
|
expected_start += int(tensor.shape[0])
|
|
|
|
if tensors:
|
|
self._write_deferred_mm_tensor(req_idx, input_batch, mm_key, tensors, first_start)
|
|
|
|
def _write_deferred_mm_tensor(
|
|
self,
|
|
req_idx: int,
|
|
input_batch: InputBatch,
|
|
mm_key: str,
|
|
tensors: list[torch.Tensor],
|
|
token_start: int,
|
|
) -> None:
|
|
mm_state = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0)
|
|
if mm_state.ndim < 2 or mm_state.shape[0] == 0:
|
|
return
|
|
|
|
if mm_key not in self.mm_cache_keys:
|
|
self.mm_outputs_cache[mm_key] = self._get_cache_tensor(
|
|
dtype=mm_state.dtype,
|
|
hidden_size=int(mm_state.shape[-1]),
|
|
)
|
|
self.mm_cache_keys.add(mm_key)
|
|
logger.info(
|
|
"Initializing deferred multimodal output cache of size %s for key: %s",
|
|
list(self.mm_outputs_cache[mm_key].shape),
|
|
mm_key,
|
|
)
|
|
|
|
slots = self._get_slot_ids_for_token_range(
|
|
req_idx=req_idx,
|
|
input_batch=input_batch,
|
|
token_start=token_start,
|
|
num_tokens=int(mm_state.shape[0]),
|
|
)
|
|
if slots.numel() == 0:
|
|
return
|
|
|
|
mm_state_cpu = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state[: slots.numel()])
|
|
mm_cache = self.mm_outputs_cache[mm_key]
|
|
flat_cache = mm_cache.view(-1, mm_cache.shape[-1])
|
|
flat_cache[slots] = mm_state_cpu
|
|
|
|
def _get_slot_ids_for_token_range(
|
|
self,
|
|
req_idx: int,
|
|
input_batch: InputBatch,
|
|
token_start: int,
|
|
num_tokens: int,
|
|
) -> torch.Tensor:
|
|
if num_tokens <= 0:
|
|
return torch.empty((0,), dtype=torch.long)
|
|
|
|
block_table = input_batch.block_table[0].block_table.cpu
|
|
token_positions = torch.arange(token_start, token_start + num_tokens, dtype=torch.long)
|
|
block_offsets = token_positions // self.block_size
|
|
max_blocks = int(block_table.shape[1])
|
|
valid = block_offsets < max_blocks
|
|
if not bool(valid.all()):
|
|
token_positions = token_positions[valid]
|
|
block_offsets = block_offsets[valid]
|
|
if token_positions.numel() == 0:
|
|
return torch.empty((0,), dtype=torch.long)
|
|
|
|
block_ids = block_table[req_idx, block_offsets].to(torch.long)
|
|
return block_ids * self.block_size + (token_positions % self.block_size)
|
|
|
|
def _coerce_to_payload_dict(
|
|
self,
|
|
element: object,
|
|
query_start_loc: torch.Tensor,
|
|
input_batch: InputBatch,
|
|
num_scheduled_tokens: dict[str, int],
|
|
) -> dict[str, object]:
|
|
"""Build the multimodal passthrough data per request for
|
|
the object under consideration. This mirrors the no-prefix-cache
|
|
path: a tensor whose first dimension matches the total scheduled
|
|
token count is sliced per request, so 1D (seq_len,) metadata that
|
|
is intentionally not cached (e.g. ref_code_len, codec_streaming)
|
|
is still split per request instead of leaking the whole batch.
|
|
"""
|
|
total_scheduled_tokens = sum(int(num_scheduled_tokens[r]) for r in input_batch.req_ids)
|
|
elem_dict = {}
|
|
for req_id in input_batch.req_ids:
|
|
req_idx = input_batch.req_id_to_index[req_id]
|
|
start = query_start_loc[req_idx]
|
|
end = start + num_scheduled_tokens[req_id]
|
|
elem_dict[req_id] = to_payload_element(
|
|
element,
|
|
req_idx,
|
|
start=start,
|
|
end=end,
|
|
pass_lists_through=True,
|
|
seq_len=total_scheduled_tokens,
|
|
)
|
|
return elem_dict
|
|
|
|
def get_merged_multimodal_states(
|
|
self,
|
|
query_start_loc: torch.Tensor,
|
|
input_batch: InputBatch,
|
|
multimodal_outputs: dict,
|
|
num_scheduled_tokens: dict[str, int],
|
|
):
|
|
"""Get the merged multimodal states if hidden state prefix caching is enabled."""
|
|
combined_multimodal_outputs = {}
|
|
# First get the prefix cached tensors that are present in the mm data
|
|
for mm_key in self.mm_cache_keys:
|
|
if mm_key in multimodal_outputs:
|
|
combined_multimodal_outputs[mm_key] = self._get_merged_tensors(
|
|
query_start_loc=query_start_loc,
|
|
input_batch=input_batch,
|
|
cache=self.mm_outputs_cache[mm_key],
|
|
hidden_states=multimodal_outputs[mm_key],
|
|
num_scheduled_tokens=num_scheduled_tokens,
|
|
)
|
|
|
|
# Then, get everything else (passthrough data); first, convert to CPU
|
|
# tensors similarly to the non prefix cached path, and then populate
|
|
# the subdicts mapping request IDs -> payload objects
|
|
passthrough_keys = set(multimodal_outputs.keys()) - self.mm_cache_keys
|
|
passthrough_mm_data = {k: v for k, v in multimodal_outputs.items() if k in passthrough_keys}
|
|
mm_cpu = build_mm_cpu(multimodal_outputs=passthrough_mm_data)
|
|
|
|
for mm_key, mm_val in mm_cpu.items():
|
|
combined_multimodal_outputs[mm_key] = self._coerce_to_payload_dict(
|
|
element=mm_val,
|
|
query_start_loc=query_start_loc,
|
|
input_batch=input_batch,
|
|
num_scheduled_tokens=num_scheduled_tokens,
|
|
)
|
|
return combined_multimodal_outputs
|
|
|
|
def get_merged_hidden_states(self, *args, **kwargs) -> dict[str, torch.Tensor]:
|
|
"""Get the merged hidden states."""
|
|
return self._get_merged_tensors(
|
|
*args,
|
|
**kwargs,
|
|
cache=self.hidden_states_cache,
|
|
)
|
|
|
|
def _get_merged_tensors(
|
|
self,
|
|
query_start_loc: torch.Tensor,
|
|
input_batch: InputBatch,
|
|
cache: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
num_scheduled_tokens: dict[str, int],
|
|
) -> dict[str, torch.Tensor]:
|
|
"""When hidden state caching is enabled, takes the input hidden_states,
|
|
which only correspond to the scheduled tokens, and returns a mapping
|
|
from request IDs to their full hidden states. This is accomplished by
|
|
looking up the block IDs & scheduled token counts to split the
|
|
hidden_states.
|
|
"""
|
|
# We do not support hybrid caches at the moment.
|
|
if len(input_batch.block_table.block_tables) > 1:
|
|
logger.warning_once(
|
|
"Omni prefix caching is enabled, but the batch block table appears to"
|
|
" have multiple kv groups; only the first group will be used!"
|
|
)
|
|
|
|
combined_hidden_states = {}
|
|
hidden_states = OmniTensorPrefixCache._coerce_to_cpu_tensor(hidden_states)
|
|
for req_id in input_batch.req_ids:
|
|
req_idx = input_batch.req_id_to_index[req_id]
|
|
|
|
if req_id in self._new_req_cache_hit_ids:
|
|
block_ids = self._get_cached_block_ids(req_idx, input_batch)
|
|
cached_hs = cache[block_ids].reshape(-1, cache.shape[-1])
|
|
|
|
# Slice the hidden states corresponding to this request;
|
|
# we do this by using the query start
|
|
start = query_start_loc[req_idx]
|
|
new_hs = hidden_states[start : start + num_scheduled_tokens[req_id]]
|
|
combined_hidden_states[req_id] = torch.cat([cached_hs, new_hs], dim=0)
|
|
else:
|
|
# cache miss for this request, pass through normally
|
|
start = query_start_loc[req_idx]
|
|
new_hs = hidden_states[start : start + num_scheduled_tokens[req_id]]
|
|
combined_hidden_states[req_id] = new_hs
|
|
|
|
return combined_hidden_states
|
|
|
|
def _get_cached_block_ids(self, req_idx: int, input_batch: InputBatch) -> torch.Tensor:
|
|
"""Given an input batch and request index in the batch (not ID), get the
|
|
block IDs corresponding to the cache hit.
|
|
"""
|
|
num_computed = input_batch.num_computed_tokens_cpu[req_idx]
|
|
# NOTE: vLLM only caches full blocks
|
|
num_cached_blocks = num_computed // self.block_size
|
|
# Get the block IDs attached to this cache hit and reindex into
|
|
# the flattened cached hidden states (i.e., 1 row per token).
|
|
return input_batch.block_table[0].block_table.cpu[req_idx, :num_cached_blocks]
|