wehub-resource-sync
378 lines
13 KiB
Python
378 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""TensorRT-LLM KV Cache Connector adapter for LMCache (in-process mode).
|
|
|
|
Implements ``LMCacheKvConnectorScheduler`` and ``LMCacheKvConnectorWorker`` —
|
|
the two classes TRT-LLM's ``kv_connector_config`` requires — backed by
|
|
an in-process LMCache engine singleton.
|
|
|
|
Lifecycle (per TRT-LLM connector ABC):
|
|
* scheduler.get_num_new_matched_tokens → engine.lookup(tokens)
|
|
* scheduler.build_connector_meta → LMCacheConnectorMetadata(loads, saves)
|
|
* worker.register_kv_caches → builds engine via _get_or_create_engine,
|
|
calls gpu_connector.register_kv_caches(kv_cache_tensor)
|
|
* worker.start_load_kv → engine.retrieve(tokens, block_ids)
|
|
* worker.wait_for_save → engine.store(tokens, block_ids)
|
|
"""
|
|
|
|
# Standard
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Optional, Tuple
|
|
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
|
|
|
|
# First Party
|
|
from lmcache import torch_dev
|
|
from lmcache.integration.tensorrt_llm.utils import (
|
|
ENGINE_NAME,
|
|
create_trtllm_metadata,
|
|
lmcache_get_config,
|
|
)
|
|
from lmcache.logging import init_logger
|
|
from lmcache.utils import EngineType, mock_up_broadcast_fn, mock_up_broadcast_object_fn
|
|
from lmcache.v1.cache_engine import LMCacheEngine, LMCacheEngineBuilder
|
|
from lmcache.v1.gpu_connector import CreateGPUConnector
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _get_or_create_engine(
|
|
llm_args: TorchLlmArgs,
|
|
kv_cache_tensor: torch.Tensor,
|
|
) -> LMCacheEngine:
|
|
"""Return the LMCache engine singleton, creating it on first call.
|
|
|
|
Called by the worker's ``register_kv_caches``. On subsequent calls
|
|
(e.g. after a model reload) the existing engine is reused and the
|
|
GPU connector is re-registered with the new KV tensor.
|
|
"""
|
|
existing = LMCacheEngineBuilder.get(ENGINE_NAME)
|
|
if existing is not None:
|
|
gpu_connector = existing.gpu_connector
|
|
if gpu_connector is not None:
|
|
gpu_connector.register_kv_caches(kv_cache_tensor) # type: ignore[attr-defined]
|
|
logger.info("LMCache TRT-LLM: reusing existing engine")
|
|
return existing
|
|
|
|
config = lmcache_get_config()
|
|
|
|
# Third Party
|
|
from transformers import AutoConfig
|
|
|
|
hf_config = AutoConfig.from_pretrained(str(llm_args.model))
|
|
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
|
|
)
|
|
num_kv_heads = num_kv_heads // llm_args.tensor_parallel_size
|
|
|
|
metadata = create_trtllm_metadata(
|
|
llm_args,
|
|
kv_cache_tensor,
|
|
config,
|
|
num_kv_heads=num_kv_heads,
|
|
head_dim=head_dim,
|
|
)
|
|
gpu_connector = CreateGPUConnector(config, metadata, EngineType.TRTLLM)
|
|
gpu_connector.register_kv_caches(kv_cache_tensor) # type: ignore[attr-defined]
|
|
|
|
engine = LMCacheEngineBuilder.get_or_create(
|
|
ENGINE_NAME,
|
|
config,
|
|
metadata,
|
|
gpu_connector,
|
|
mock_up_broadcast_fn,
|
|
mock_up_broadcast_object_fn,
|
|
)
|
|
engine.post_init()
|
|
|
|
logger.info(
|
|
"LMCache TRT-LLM: created engine (chunk_size=%d, tensor_shape=%s, dtype=%s)",
|
|
config.chunk_size,
|
|
list(kv_cache_tensor.shape),
|
|
kv_cache_tensor.dtype,
|
|
)
|
|
return engine
|
|
|
|
|
|
def destroy_engine() -> None:
|
|
"""Destroy the engine singleton. Safe to call if not created."""
|
|
if LMCacheEngineBuilder.get(ENGINE_NAME) is not None:
|
|
LMCacheEngineBuilder.destroy(ENGINE_NAME)
|
|
logger.info("LMCache TRT-LLM: engine destroyed")
|
|
|
|
|
|
@dataclass
|
|
class _BlockSpec:
|
|
"""Tokens and block IDs for a single load or store operation."""
|
|
|
|
tokens: List[int]
|
|
block_ids: List[int]
|
|
|
|
|
|
@dataclass
|
|
class LMCacheConnectorMetadata:
|
|
"""Metadata passed from the scheduler to the worker each forward step."""
|
|
|
|
loads: dict = field(default_factory=dict)
|
|
saves: dict = field(default_factory=dict)
|
|
|
|
|
|
class LMCacheKvConnectorScheduler(KvCacheConnectorScheduler):
|
|
"""Scheduler-side connector hook.
|
|
|
|
Queries the LMCache engine for cached token counts and emits
|
|
per-request load/save block specs for the worker to act on.
|
|
"""
|
|
|
|
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) — set by
|
|
# get_num_new_matched_tokens, consumed by build_connector_meta.
|
|
self._pending: dict = {}
|
|
# Engine is created by the worker in register_kv_caches, which
|
|
# may run concurrently with scheduler init. Resolved lazily.
|
|
self._engine: Optional[LMCacheEngine] = None
|
|
|
|
def get_num_new_matched_tokens(
|
|
self,
|
|
request: LlmRequest,
|
|
num_computed_tokens: int,
|
|
) -> Tuple[int, bool]:
|
|
"""Return how many additional tokens LMCache can provide beyond
|
|
``num_computed_tokens`` (which TRT-LLM matched via GPU block reuse).
|
|
|
|
Args:
|
|
request: The incoming request with its full token sequence.
|
|
num_computed_tokens: Tokens already matched on device
|
|
(block-aligned).
|
|
|
|
Returns:
|
|
``(new_matched, is_async)``. ``is_async`` is always
|
|
``False`` in this adapter.
|
|
"""
|
|
t0 = time.perf_counter()
|
|
|
|
if not self._engine:
|
|
self._engine = LMCacheEngineBuilder.get(ENGINE_NAME)
|
|
if not self._engine:
|
|
self._pending[request.request_id] = ([], 0)
|
|
return 0, False
|
|
|
|
# TRT-LLM should always pass block-aligned positions.
|
|
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)
|
|
logger.debug(
|
|
"LMCache TRT-LLM scheduler: req %d short-circuit "
|
|
"(TRT matched %d of %d block-aligned tokens) %.3fms",
|
|
request.request_id,
|
|
num_computed_tokens,
|
|
max_block_aligned,
|
|
(time.perf_counter() - t0) * 1000,
|
|
)
|
|
return 0, False
|
|
|
|
t1 = time.perf_counter()
|
|
cached_tokens = self._engine.lookup(tokens=all_tokens)
|
|
t2 = time.perf_counter()
|
|
|
|
new_matched = max(0, cached_tokens - num_computed_tokens)
|
|
new_matched = (new_matched // self._block_size) * self._block_size
|
|
|
|
self._pending[request.request_id] = (all_tokens, new_matched)
|
|
|
|
logger.debug(
|
|
"LMCache TRT-LLM 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
|
|
) -> LMCacheConnectorMetadata:
|
|
"""Build per-request load/save specs from the pending lookup
|
|
results. The runtime binds the returned metadata to the worker
|
|
via ``bind_connector_meta`` before the forward pass starts.
|
|
"""
|
|
meta = LMCacheConnectorMetadata()
|
|
|
|
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:
|
|
"""Return whether async saving is in progress.
|
|
|
|
Always ``False`` — saves are synchronous in this adapter.
|
|
"""
|
|
return False
|
|
|
|
def update_state_after_alloc(
|
|
self, request: LlmRequest, block_ids: List[int]
|
|
) -> None:
|
|
"""No-op — block IDs are captured in ``build_connector_meta``
|
|
from ``scheduler_output.new_requests``.
|
|
"""
|
|
pass
|
|
|
|
|
|
class LMCacheKvConnectorWorker(KvCacheConnectorWorker):
|
|
"""Worker-side connector hook.
|
|
|
|
Performs GPU↔CPU KV transfers via the LMCache engine.
|
|
"""
|
|
|
|
def __init__(self, llm_args: TorchLlmArgs) -> None:
|
|
super().__init__(llm_args)
|
|
self._block_size: int = self._llm_args.kv_cache_config.tokens_per_block
|
|
# Cached after register_kv_caches to avoid per-call singleton lookup.
|
|
self._engine: Optional[LMCacheEngine] = None
|
|
self._load_stream: Optional[torch_dev.Stream] = None
|
|
self._store_stream: Optional[torch_dev.Stream] = None
|
|
|
|
@property
|
|
def _meta(self) -> Optional[LMCacheConnectorMetadata]:
|
|
"""Typed accessor for ``self._metadata`` set by the base class."""
|
|
return self._metadata # type: ignore[return-value]
|
|
|
|
def register_kv_caches(self, kv_cache_tensor: torch.Tensor) -> None:
|
|
"""Register the KV cache tensor and create the LMCache engine.
|
|
|
|
Called once by the runtime after KV cache allocation. Caches the
|
|
engine and its load/store streams for fast access on every step.
|
|
"""
|
|
self._engine = _get_or_create_engine(
|
|
llm_args=self._llm_args,
|
|
kv_cache_tensor=kv_cache_tensor,
|
|
)
|
|
gpu_conn = self._engine.gpu_connector
|
|
if gpu_conn is not None:
|
|
self._load_stream = gpu_conn.load_stream # type: ignore[attr-defined]
|
|
self._store_stream = gpu_conn.store_stream # type: ignore[attr-defined]
|
|
|
|
def start_load_kv(self, stream: torch_dev.Stream) -> None:
|
|
"""Load KV blocks from LMCache into the GPU paged cache.
|
|
|
|
Retrieves all pending blocks on the load stream, then syncs the
|
|
forward-pass stream against it. The cross-layer format loads
|
|
every layer in a single kernel — no per-layer overlap to exploit.
|
|
"""
|
|
meta = self._meta
|
|
if meta is None or not meta.loads or self._engine is None:
|
|
return
|
|
|
|
t0 = time.perf_counter()
|
|
for spec in meta.loads.values():
|
|
if not spec.tokens or not spec.block_ids:
|
|
continue
|
|
self._engine.retrieve(tokens=spec.tokens, block_ids=spec.block_ids)
|
|
|
|
if self._load_stream is not None:
|
|
stream.wait_stream(self._load_stream)
|
|
|
|
logger.debug(
|
|
"LMCache TRT-LLM 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 — cross-layer loads complete in :meth:`start_load_kv`."""
|
|
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:
|
|
"""Store newly computed KV blocks from GPU to LMCache's CPU cache.
|
|
|
|
Waits on the forward-pass stream, runs ``engine.store`` for each
|
|
request with new blocks, and synchronizes the store stream
|
|
before returning.
|
|
"""
|
|
meta = self._meta
|
|
if meta is None or not meta.saves or self._engine is None:
|
|
return
|
|
|
|
t0 = time.perf_counter()
|
|
if self._store_stream is not None:
|
|
self._store_stream.wait_stream(stream)
|
|
t1 = time.perf_counter()
|
|
|
|
for spec in meta.saves.values():
|
|
if not spec.tokens or not spec.block_ids:
|
|
continue
|
|
self._engine.store(tokens=spec.tokens, block_ids=spec.block_ids)
|
|
t2 = time.perf_counter()
|
|
|
|
if self._store_stream is not None:
|
|
self._store_stream.synchronize()
|
|
t3 = time.perf_counter()
|
|
|
|
logger.debug(
|
|
"LMCache TRT-LLM worker: wait_for_save stream_wait=%.3fms "
|
|
"store=%.3fms sync=%.3fms num_saves=%d",
|
|
(t1 - t0) * 1000,
|
|
(t2 - t1) * 1000,
|
|
(t3 - t2) * 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 ops are synchronous here — nothing is ever pending."""
|
|
return [], []
|