wehub-resource-sync
76 lines
2.7 KiB
Python
76 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""EngineDrivenContext to store/retrieve a contiguous KV tensor for SDK use."""
|
|
|
|
# Future
|
|
from __future__ import annotations
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.v1.multiprocess.custom_types import IPCCacheServerKey
|
|
from lmcache.v1.multiprocess.transfer_context.base import EngineDrivenContext
|
|
|
|
|
|
class ContiguousTransferWrapper:
|
|
"""Store/retrieve a contiguous KV tensor through an ``EngineDrivenContext``.
|
|
|
|
Args:
|
|
context: The engine-driven (SHM or pickle) transport.
|
|
chunk_size: Number of tokens per LMCache chunk.
|
|
"""
|
|
|
|
def __init__(self, context: EngineDrivenContext, chunk_size: int) -> None:
|
|
self._context = context
|
|
self._chunk_size = chunk_size
|
|
|
|
def store(self, key: IPCCacheServerKey, instance_id: int, kv: torch.Tensor) -> bool:
|
|
"""Store a contiguous [2, L, T, D] tensor
|
|
|
|
Args:
|
|
key: The cache server key.
|
|
instance_id: The cache server instance ID.
|
|
kv: The contiguous KV tensor to store.
|
|
|
|
Returns:
|
|
True if the store was successful, False otherwise.
|
|
"""
|
|
result = self._context.prepare_store(key, instance_id)
|
|
if result is None:
|
|
# Pickle: chunk the contiguous KV tensor (commit takes list of chunks).
|
|
num_chunks = kv.shape[2] // self._chunk_size
|
|
chunks = [
|
|
kv[
|
|
:, :, i * self._chunk_size : (i + 1) * self._chunk_size, :
|
|
].contiguous()
|
|
for i in range(num_chunks)
|
|
]
|
|
else:
|
|
# SHM: fill missing chunks' slots in place.
|
|
slot_tensors, chunk_indices = result
|
|
for slot, chunk_idx in zip(slot_tensors, chunk_indices, strict=True):
|
|
start = chunk_idx * self._chunk_size
|
|
slot.copy_(kv[:, :, start : start + self._chunk_size, :])
|
|
chunks = []
|
|
return self._context.commit_store(key, instance_id, chunks)
|
|
|
|
def retrieve(self, key: IPCCacheServerKey, instance_id: int) -> torch.Tensor | None:
|
|
"""Retrieve the KV as a contiguous [2, L, hit_tokens, D] tensor
|
|
|
|
Args:
|
|
key: The cache server key.
|
|
instance_id: The cache server instance ID.
|
|
|
|
Returns:
|
|
The contiguous KV tensor if found, None otherwise.
|
|
"""
|
|
slot_tensors = self._context.prepare_retrieve(key, instance_id)
|
|
if not slot_tensors:
|
|
return None
|
|
try:
|
|
# Both Pickle and SHM returns list of [2, L, T, D] tensors
|
|
# Concatenate along the token dimension.
|
|
return torch.cat(slot_tensors, dim=2)
|
|
finally:
|
|
self._context.commit_retrieve(key, instance_id)
|