wehub-resource-sync
489 lines
17 KiB
Python
489 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""TensorRT-LLM KV Cache Connector adapter for LMCache (multi-process mode).
|
|
|
|
Implements ``LMCacheMPKvConnectorScheduler`` and
|
|
``LMCacheMPKvConnectorWorker`` — the two classes TRT-LLM's
|
|
``kv_connector_config`` requires — backed by a standalone LMCache server
|
|
reached over ZMQ. Provides process isolation and shared caching across
|
|
multiple TRT-LLM instances on the same node.
|
|
|
|
The KV pool tensor is shared with the server via :class:`RawCudaIPCWrapper`
|
|
because TRT-LLM's pool is allocated outside PyTorch's caching allocator
|
|
(``at::for_blob`` over ``cudaMalloc``), which makes
|
|
``UntypedStorage._share_cuda_()`` raise. The wrapper bypasses that path.
|
|
"""
|
|
|
|
# Standard
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Optional, Tuple
|
|
import os
|
|
import time
|
|
|
|
# Third Party
|
|
from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import (
|
|
KvCacheConnectorScheduler,
|
|
KvCacheConnectorWorker,
|
|
SchedulerOutput,
|
|
)
|
|
from tensorrt_llm.bindings.internal.batch_manager import LlmRequest
|
|
from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
|
|
import torch
|
|
import zmq
|
|
|
|
# First Party
|
|
from lmcache import torch_dev
|
|
from lmcache.logging import init_logger
|
|
from lmcache.utils import EngineType, check_interprocess_event_support
|
|
from lmcache.v1.multiprocess.custom_types import (
|
|
IPCCacheServerKey,
|
|
)
|
|
from lmcache.v1.multiprocess.mq import MessageQueueClient, MessagingFuture
|
|
from lmcache.v1.multiprocess.protocol import RequestType, get_response_class
|
|
from lmcache.v1.platform.cuda.ipc_wrapper import RawCudaIPCWrapper
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
DEFAULT_SERVER_URL = "ipc:///tmp/lmcache.sock"
|
|
DEFAULT_MQ_TIMEOUT: float = 300.0
|
|
|
|
|
|
def _get_server_url(llm_args: "TorchLlmArgs") -> str:
|
|
"""Resolve the server URL: connector-config field > env var > default."""
|
|
cfg = llm_args.kv_connector_config
|
|
if cfg is not None and cfg.server_url is not None:
|
|
return cfg.server_url
|
|
return os.environ.get("LMCACHE_SERVER_URL", DEFAULT_SERVER_URL)
|
|
|
|
|
|
def _send_request(
|
|
mq_client: MessageQueueClient,
|
|
request_type: RequestType,
|
|
payloads: list,
|
|
) -> MessagingFuture:
|
|
return mq_client.submit_request(
|
|
request_type, payloads, get_response_class(request_type)
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _BlockSpec:
|
|
tokens: List[int]
|
|
block_ids: List[int]
|
|
|
|
|
|
@dataclass
|
|
class LMCacheMPConnectorMetadata:
|
|
loads: dict = field(default_factory=dict)
|
|
saves: dict = field(default_factory=dict)
|
|
|
|
|
|
class LMCacheMPKvConnectorScheduler(KvCacheConnectorScheduler):
|
|
"""TRT-LLM scheduler that routes lookup requests to an LMCache MP server."""
|
|
|
|
def __init__(self, llm_args: TorchLlmArgs) -> None:
|
|
super().__init__(llm_args)
|
|
self._block_size: int = self._llm_args.kv_cache_config.tokens_per_block
|
|
# request_id -> (all_tokens, num_matched).
|
|
self._pending: dict = {}
|
|
|
|
self._zmq_context = zmq.Context()
|
|
self._mq_client = MessageQueueClient(
|
|
_get_server_url(self._llm_args), self._zmq_context
|
|
)
|
|
self._mq_timeout = float(
|
|
os.environ.get("LMCACHE_MQ_TIMEOUT", DEFAULT_MQ_TIMEOUT)
|
|
)
|
|
|
|
future = _send_request(self._mq_client, RequestType.GET_CHUNK_SIZE, [])
|
|
self._chunk_size = future.result(timeout=self._mq_timeout)
|
|
logger.info(
|
|
"LMCache MP scheduler: connected to server at %s (chunk_size=%d)",
|
|
_get_server_url(self._llm_args),
|
|
self._chunk_size,
|
|
)
|
|
|
|
# Third Party
|
|
import tensorrt_llm
|
|
|
|
self._rank = tensorrt_llm.mpi_rank()
|
|
tp_size = llm_args.tensor_parallel_size
|
|
pp_size = llm_args.pipeline_parallel_size
|
|
self._world_size = tp_size * pp_size
|
|
self._model_name = str(getattr(llm_args, "model", "unknown_model"))
|
|
|
|
def _create_key(
|
|
self,
|
|
token_ids: List[int],
|
|
start: int,
|
|
end: int,
|
|
request_id: int,
|
|
) -> IPCCacheServerKey:
|
|
return IPCCacheServerKey(
|
|
model_name=self._model_name,
|
|
world_size=self._world_size,
|
|
worker_id=None,
|
|
token_ids=tuple(token_ids),
|
|
start=start,
|
|
end=end,
|
|
request_id=str(request_id),
|
|
)
|
|
|
|
def get_num_new_matched_tokens(
|
|
self,
|
|
request: LlmRequest,
|
|
num_computed_tokens: int,
|
|
) -> Tuple[int, bool]:
|
|
"""Return how many additional tokens the LMCache server can provide.
|
|
|
|
Submits a ``LOOKUP`` to the server and queries
|
|
``QUERY_PREFETCH_STATUS`` by ``request_id`` to read the result.
|
|
``LOOKUP`` returns ``None`` on the server protocol — the prefetch
|
|
is tracked server-side keyed by ``request_id``.
|
|
"""
|
|
t0 = time.perf_counter()
|
|
|
|
if num_computed_tokens % self._block_size != 0:
|
|
self._pending[request.request_id] = ([], 0)
|
|
return 0, False
|
|
|
|
all_tokens = list(request.get_tokens(0))
|
|
|
|
max_block_aligned = (len(all_tokens) // self._block_size) * self._block_size
|
|
if num_computed_tokens >= max_block_aligned:
|
|
self._pending[request.request_id] = (all_tokens, 0)
|
|
return 0, False
|
|
|
|
aligned_end = (len(all_tokens) // self._chunk_size) * self._chunk_size
|
|
key = self._create_key(
|
|
all_tokens, start=0, end=aligned_end, request_id=request.request_id
|
|
).no_worker_id_version()
|
|
|
|
t1 = time.perf_counter()
|
|
|
|
try:
|
|
_send_request(self._mq_client, RequestType.LOOKUP, [key, 1]).result(
|
|
timeout=self._mq_timeout
|
|
)
|
|
result = _send_request(
|
|
self._mq_client,
|
|
RequestType.QUERY_PREFETCH_STATUS,
|
|
[str(request.request_id)],
|
|
).result(timeout=self._mq_timeout)
|
|
cached_tokens = result * self._chunk_size if result is not None else 0
|
|
except Exception as e:
|
|
logger.warning("LMCache MP scheduler: lookup failed: %s", e)
|
|
self._pending[request.request_id] = (all_tokens, 0)
|
|
return 0, False
|
|
|
|
t2 = time.perf_counter()
|
|
|
|
new_matched = max(0, cached_tokens - num_computed_tokens)
|
|
new_matched = (new_matched // self._block_size) * self._block_size
|
|
|
|
# ``LOOKUP`` acquires read locks on chunks in [0, cached_tokens).
|
|
# TRT-LLM already has [0, num_computed_tokens), so those chunks
|
|
# will never be retrieved — release their locks (chunk-aligned)
|
|
# to avoid holding them until TTL expiry.
|
|
overlap_end = min(cached_tokens, num_computed_tokens)
|
|
overlap_end = (overlap_end // self._chunk_size) * self._chunk_size
|
|
if overlap_end > 0:
|
|
free_key = self._create_key(
|
|
all_tokens,
|
|
start=0,
|
|
end=overlap_end,
|
|
request_id=request.request_id,
|
|
).no_worker_id_version()
|
|
try:
|
|
_send_request(
|
|
self._mq_client,
|
|
RequestType.FREE_LOOKUP_LOCKS,
|
|
[free_key, 1],
|
|
)
|
|
except Exception as e:
|
|
logger.warning("LMCache MP scheduler: free_lookup_locks failed: %s", e)
|
|
|
|
self._pending[request.request_id] = (all_tokens, new_matched)
|
|
|
|
logger.debug(
|
|
"LMCache MP scheduler: req %d lookup=%.3fms total=%.3fms "
|
|
"trt_matched=%d lmcache_cached=%d new_matched=%d",
|
|
request.request_id,
|
|
(t2 - t1) * 1000,
|
|
(time.perf_counter() - t0) * 1000,
|
|
num_computed_tokens,
|
|
cached_tokens,
|
|
new_matched,
|
|
)
|
|
return new_matched, False
|
|
|
|
def build_connector_meta(
|
|
self, scheduler_output: SchedulerOutput
|
|
) -> LMCacheMPConnectorMetadata:
|
|
"""Build per-request load/save specs from pending lookup results."""
|
|
meta = LMCacheMPConnectorMetadata()
|
|
|
|
for req in scheduler_output.new_requests:
|
|
if req.request_id not in self._pending:
|
|
continue
|
|
|
|
all_tokens, num_matched = self._pending[req.request_id]
|
|
block_ids: List[int] = list(req.new_block_ids)
|
|
num_computed_blocks = req.computed_position // self._block_size
|
|
|
|
if num_matched > 0:
|
|
meta.loads[req.request_id] = _BlockSpec(
|
|
tokens=all_tokens, block_ids=block_ids
|
|
)
|
|
|
|
save_start = max(num_computed_blocks, num_matched // self._block_size)
|
|
num_full_new_blocks = len(req.new_tokens) // self._block_size
|
|
if (
|
|
save_start < len(block_ids)
|
|
and num_full_new_blocks > 0
|
|
and save_start < num_computed_blocks + num_full_new_blocks
|
|
):
|
|
meta.saves[req.request_id] = _BlockSpec(
|
|
tokens=all_tokens, block_ids=block_ids
|
|
)
|
|
|
|
self._pending.clear()
|
|
return meta
|
|
|
|
def request_finished(self, request: LlmRequest, cache_block_ids: List[int]) -> bool:
|
|
"""Notify the server so it can clean up per-request state.
|
|
|
|
Saves are synchronous in this adapter, so we never need to defer
|
|
deallocation — return ``False``. We still call ``END_SESSION`` to
|
|
release the server-side token-hash/session state for the request.
|
|
"""
|
|
try:
|
|
_send_request(
|
|
self._mq_client,
|
|
RequestType.END_SESSION,
|
|
[str(request.request_id)],
|
|
)
|
|
except Exception as e:
|
|
logger.warning("LMCache MP scheduler: end_session failed: %s", e)
|
|
return False
|
|
|
|
def update_state_after_alloc(
|
|
self, request: LlmRequest, block_ids: List[int]
|
|
) -> None:
|
|
"""No-op — block IDs are captured in :meth:`build_connector_meta`."""
|
|
pass
|
|
|
|
|
|
class LMCacheMPKvConnectorWorker(KvCacheConnectorWorker):
|
|
"""TRT-LLM worker that routes store/retrieve to an LMCache MP server."""
|
|
|
|
def __init__(self, llm_args: TorchLlmArgs) -> None:
|
|
super().__init__(llm_args)
|
|
self._block_size: int = self._llm_args.kv_cache_config.tokens_per_block
|
|
|
|
self._zmq_context = zmq.Context()
|
|
self._mq_client = MessageQueueClient(
|
|
_get_server_url(self._llm_args), self._zmq_context
|
|
)
|
|
self._mq_timeout = float(
|
|
os.environ.get("LMCACHE_MQ_TIMEOUT", DEFAULT_MQ_TIMEOUT)
|
|
)
|
|
|
|
self._instance_id = os.getpid()
|
|
self._registered = False
|
|
|
|
# Third Party
|
|
import tensorrt_llm
|
|
|
|
self._rank = tensorrt_llm.mpi_rank()
|
|
tp_size = llm_args.tensor_parallel_size
|
|
pp_size = llm_args.pipeline_parallel_size
|
|
self._world_size = tp_size * pp_size
|
|
self._model_name = str(getattr(llm_args, "model", "unknown_model"))
|
|
|
|
future = _send_request(self._mq_client, RequestType.GET_CHUNK_SIZE, [])
|
|
self._chunk_size = future.result(timeout=self._mq_timeout)
|
|
|
|
def _create_key(
|
|
self,
|
|
token_ids: List[int],
|
|
request_id: int,
|
|
) -> IPCCacheServerKey:
|
|
aligned_end = (len(token_ids) // self._chunk_size) * self._chunk_size
|
|
return IPCCacheServerKey(
|
|
model_name=self._model_name,
|
|
world_size=self._world_size,
|
|
worker_id=self._rank,
|
|
token_ids=tuple(token_ids),
|
|
start=0,
|
|
end=aligned_end,
|
|
request_id=str(request_id),
|
|
)
|
|
|
|
def register_kv_caches(self, kv_cache_tensor: torch.Tensor) -> None:
|
|
"""Register the KV pool with the LMCache server via raw CUDA IPC.
|
|
|
|
TRT-LLM provides a 4-D pool tensor
|
|
``[NB, NL, 2, NH * BS * HS]``. The server reshapes it to 6-D
|
|
``[NB, NL, 2, NH, BS, HS]`` from the ``layout_hints`` so format
|
|
detection lands on ``NB_NL_TWO_NH_BS_HS``.
|
|
"""
|
|
if self._registered:
|
|
logger.info("LMCache MP worker: KV caches already registered")
|
|
return
|
|
|
|
# Third Party
|
|
from transformers import AutoConfig
|
|
|
|
hf_config = AutoConfig.from_pretrained(self._model_name)
|
|
head_dim = getattr(
|
|
hf_config,
|
|
"head_dim",
|
|
hf_config.hidden_size // hf_config.num_attention_heads,
|
|
)
|
|
num_kv_heads = getattr(
|
|
hf_config, "num_key_value_heads", hf_config.num_attention_heads
|
|
)
|
|
tp_size = self._llm_args.tensor_parallel_size
|
|
num_kv_heads = num_kv_heads // tp_size
|
|
|
|
_, _, _, block_size_flat = kv_cache_tensor.shape
|
|
tokens_per_block = block_size_flat // (num_kv_heads * head_dim)
|
|
|
|
wrapped = [RawCudaIPCWrapper(kv_cache_tensor)]
|
|
|
|
layout_hints = {
|
|
"kv_layout": "HND",
|
|
"num_kv_heads": num_kv_heads,
|
|
"tokens_per_block": tokens_per_block,
|
|
"head_dim": head_dim,
|
|
}
|
|
|
|
future = _send_request(
|
|
self._mq_client,
|
|
RequestType.REGISTER_KV_CACHE,
|
|
[
|
|
self._instance_id,
|
|
wrapped,
|
|
self._model_name,
|
|
self._world_size,
|
|
EngineType.TRTLLM,
|
|
layout_hints,
|
|
[],
|
|
],
|
|
)
|
|
try:
|
|
future.result(timeout=self._mq_timeout)
|
|
self._registered = True
|
|
logger.info(
|
|
"LMCache MP worker: registered KV caches "
|
|
"(tensor_shape=%s, NH=%d, BS=%d, HS=%d)",
|
|
list(kv_cache_tensor.shape),
|
|
num_kv_heads,
|
|
tokens_per_block,
|
|
head_dim,
|
|
)
|
|
except TimeoutError:
|
|
logger.error(
|
|
"LMCache MP worker: KV cache registration timed out after %ss",
|
|
self._mq_timeout,
|
|
)
|
|
|
|
def start_load_kv(self, stream: torch_dev.Stream) -> None:
|
|
"""Send ``RETRIEVE`` requests for each pending load."""
|
|
meta: Optional[LMCacheMPConnectorMetadata] = self._metadata
|
|
if meta is None or not meta.loads:
|
|
return
|
|
|
|
t0 = time.perf_counter()
|
|
# Not all backends support interprocess Events (CUDA IPC specific)
|
|
check_interprocess_event_support()
|
|
event = torch_dev.Event(interprocess=True)
|
|
event.record(stream)
|
|
|
|
for req_id, spec in meta.loads.items():
|
|
if not spec.tokens or not spec.block_ids:
|
|
continue
|
|
|
|
key = self._create_key(spec.tokens, req_id)
|
|
try:
|
|
_send_request(
|
|
self._mq_client,
|
|
RequestType.RETRIEVE,
|
|
[
|
|
key,
|
|
self._instance_id,
|
|
[spec.block_ids],
|
|
event.ipc_handle(),
|
|
0, # skip_first_n_tokens
|
|
],
|
|
).result(timeout=self._mq_timeout)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"LMCache MP worker: retrieve failed for req %d: %s",
|
|
req_id,
|
|
e,
|
|
)
|
|
|
|
logger.debug(
|
|
"LMCache MP worker: start_load_kv retrieve=%.3fms num_loads=%d",
|
|
(time.perf_counter() - t0) * 1000,
|
|
len(meta.loads),
|
|
)
|
|
|
|
def wait_for_layer_load(self, layer_idx: int, stream: torch_dev.Stream) -> None:
|
|
"""No-op — server synchronizes via CUDA IPC events."""
|
|
pass
|
|
|
|
def save_kv_layer(self, layer_idx: int, stream: torch_dev.Stream) -> None:
|
|
"""No-op — saves are batched in :meth:`wait_for_save`."""
|
|
pass
|
|
|
|
def wait_for_save(self, stream: torch_dev.Stream) -> None:
|
|
"""Send ``STORE`` requests for each pending save."""
|
|
meta: Optional[LMCacheMPConnectorMetadata] = self._metadata
|
|
if meta is None or not meta.saves:
|
|
return
|
|
|
|
t0 = time.perf_counter()
|
|
# Not all backends support interprocess Events (CUDA IPC specific)
|
|
check_interprocess_event_support()
|
|
event = torch_dev.Event(interprocess=True)
|
|
event.record(stream)
|
|
|
|
for req_id, spec in meta.saves.items():
|
|
if not spec.tokens or not spec.block_ids:
|
|
continue
|
|
|
|
key = self._create_key(spec.tokens, req_id)
|
|
try:
|
|
_send_request(
|
|
self._mq_client,
|
|
RequestType.STORE,
|
|
[
|
|
key,
|
|
self._instance_id,
|
|
[spec.block_ids],
|
|
event.ipc_handle(),
|
|
],
|
|
).result(timeout=self._mq_timeout)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"LMCache MP worker: store failed for req %d: %s",
|
|
req_id,
|
|
e,
|
|
)
|
|
|
|
logger.debug(
|
|
"LMCache MP worker: wait_for_save store=%.3fms num_saves=%d",
|
|
(time.perf_counter() - t0) * 1000,
|
|
len(meta.saves),
|
|
)
|
|
|
|
def get_finished(
|
|
self,
|
|
finished_gen_req_ids: List[int],
|
|
started_loading_req_ids: List[int],
|
|
) -> Tuple[List[int], List[int]]:
|
|
"""All operations are synchronous — nothing is ever pending."""
|
|
return [], []
|