diff --git a/examples/llm-api/configs/trtllm_kvbm_connector_extra.yaml b/examples/llm-api/configs/trtllm_kvbm_connector_extra.yaml new file mode 100644 index 000000000000..89356108bea2 --- /dev/null +++ b/examples/llm-api/configs/trtllm_kvbm_connector_extra.yaml @@ -0,0 +1,15 @@ +# Extra LLM API options for trtllm-serve with the Dynamo KVBM connector. +# +# Prerequisites: +# - dynamo kvbm installed: see https://github.com/ai-dynamo/dynamo +# - PyTorch backend +# +# Example: +# trtllm-serve Qwen/Qwen2-1.5B-Instruct --backend pytorch --host 0.0.0.0 --port 8000 \ +# --extra_llm_api_options /path/to/this/file + +kv_cache_config: + enable_block_reuse: true + +kv_connector_config: + connector: kvbm diff --git a/examples/llm-api/configs/trtllm_lmcache_connector_extra.yaml b/examples/llm-api/configs/trtllm_lmcache_connector_extra.yaml new file mode 100644 index 000000000000..d855c361672b --- /dev/null +++ b/examples/llm-api/configs/trtllm_lmcache_connector_extra.yaml @@ -0,0 +1,20 @@ +# Extra LLM API options for trtllm-serve with the LMCache KV connector. +# +# Prerequisites: +# - lmcache installed: pip install lmcache +# - PyTorch backend (LMCache connector is wired through the PyTorch executor) +# +# IMPORTANT: Set PYTHONHASHSEED=0 before starting the process for +# deterministic cache key hashing in LMCache. +# +# Example: +# PYTHONHASHSEED=0 trtllm-serve Qwen/Qwen2-1.5B-Instruct --backend pytorch --host 0.0.0.0 --port 8000 \ +# --trust_remote_code \ +# --extra_llm_api_options /path/to/this/file + +kv_cache_config: + enable_block_reuse: true + max_tokens: 150000 + +kv_connector_config: + connector: lmcache diff --git a/examples/llm-api/llm_kv_cache_connector.py b/examples/llm-api/llm_kv_cache_connector.py index 2e87e0c68682..6b890d8ca88a 100644 --- a/examples/llm-api/llm_kv_cache_connector.py +++ b/examples/llm-api/llm_kv_cache_connector.py @@ -90,7 +90,7 @@ import torch from tensorrt_llm import LLM, SamplingParams, logger -from tensorrt_llm._torch.pyexecutor.kv_cache_connector import ( +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 KvCacheConnectorConfig, TorchLlmArgs diff --git a/examples/llm-api/llm_lmcache_connector.py b/examples/llm-api/llm_lmcache_connector.py new file mode 100644 index 000000000000..cff69192c05e --- /dev/null +++ b/examples/llm-api/llm_lmcache_connector.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +### :title LMCache KV Cache Connector +### :order 7 +### :section Customization +"""Demonstrates using LMCache as a KV cache backend for TensorRT-LLM. + +Uses the KV Cache Connector interface. + +LMCache stores previously computed KV tensors and replays them on subsequent +requests with the same prefix, reducing recomputation. + +The connector implementation lives in LMCache: + lmcache.integration.tensorrt_llm.tensorrt_adapter + +TRT-LLM resolves the ``"lmcache"`` preset to the correct import paths +automatically via the connector registry. + +Prerequisites: + pip install lmcache + +How to run: + PYTHONHASHSEED=0 python llm_lmcache_connector.py Qwen/Qwen2-1.5B-Instruct + +Note: PYTHONHASHSEED=0 must be set before the Python process starts +to ensure deterministic cache key hashing in LMCache. + +Expected output: + Second request logs show "Retrieved N tokens" and both outputs are identical. + +See Also: + examples/llm-api/configs/trtllm_lmcache_connector_extra.yaml -- trtllm-serve YAML +""" + +import click + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, KvCacheConnectorConfig + +try: + from lmcache.integration.tensorrt_llm import destroy_engine +except ImportError as e: + raise ImportError( + "LMCache is not installed or is missing the TensorRT-LLM integration. " + "Run: pip install 'lmcache'" + ) from e + +# A prompt long enough to produce at least one full TRT-LLM KV block. +_TEST_PROMPT = ( + "Nvidia Corporation is an American technology company headquartered in " + "Santa Clara, California. Founded in 1993 by Jensen Huang, Chris " + "Malachowsky, and Curtis Priem, it develops graphics processing units " + "(GPUs), system on a chips (SoCs), and application programming " + "interfaces (APIs) for data science, high-performance computing, and " + "mobile and automotive applications. Tell me about the company." +) + + +@click.command() +@click.argument("model", type=str) +def main(model: str): + kv_cache_config = KvCacheConfig(enable_block_reuse=True) + kv_connector_config = KvCacheConnectorConfig(connector="lmcache") + sampling_params = SamplingParams(max_tokens=32) + + # Both requests go to the same LLM instance so the in-process LMCache + # engine (and its CPU memory cache) survives between the two calls. + llm = LLM( + model=model, + backend="pytorch", + kv_cache_config=kv_cache_config, + kv_connector_config=kv_connector_config, + ) + + print("--- First request (cold cache, KV will be computed and stored) ---") + output0 = llm.generate([_TEST_PROMPT], sampling_params) + text0 = output0[0].outputs[0].text + print("First output:", text0) + + print("\n--- Second request (warm cache, KV should be retrieved) ---") + output1 = llm.generate([_TEST_PROMPT], sampling_params) + text1 = output1[0].outputs[0].text + print("Second output (using LMCache KV cache):", text1) + + assert text0 == text1, ( + f"Outputs differ — cache reuse may not have worked correctly.\n" + f"First: {text0!r}\n" + f"Second: {text1!r}" + ) + print("\nOK: outputs match, LMCache KV reuse confirmed.") + + destroy_engine() + + +if __name__ == "__main__": + main() diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b000427620d7..261c2a3f085c 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -31,9 +31,9 @@ get_spec_decoder, should_use_separate_draft_kv_cache) from .config_utils import (get_qwen3_hybrid_layer_masks, is_mla, is_nemotron_hybrid, is_qwen3_hybrid) +from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder -from .kv_cache_connector import KvCacheConnectorManager from .kv_cache_transceiver import AttentionTypeCpp, create_kv_cache_transceiver from .llm_request import ExecutorResponse from .mamba_cache_manager import MambaHybridCacheManager diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/__init__.py new file mode 100644 index 000000000000..0f3ff0444978 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tensorrt_llm._torch.pyexecutor.connectors.registry import CONNECTOR_REGISTRY + +__all__ = ["CONNECTOR_REGISTRY"] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py similarity index 78% rename from tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py rename to tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 5715afe96045..e56564b69656 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -37,23 +37,23 @@ from abc import ABC, abstractmethod from collections import defaultdict from dataclasses import dataclass, field -from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, - Tuple) +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple import torch from tensorrt_llm._utils import mpi_allgather, mpi_broadcast, mpi_rank from tensorrt_llm.bindings import LlmRequestState -from tensorrt_llm.bindings.internal.batch_manager import \ - KvCacheConnectorManager as KvCacheConnectorManagerCpp +from tensorrt_llm.bindings.internal.batch_manager import ( + KvCacheConnectorManager as KvCacheConnectorManagerCpp, +) from tensorrt_llm.bindings.internal.batch_manager import LlmRequest from tensorrt_llm.llmapi.llm_args import TorchLlmArgs -from .llm_request import get_draft_token_length -from .scheduler import ScheduledRequests +from ..llm_request import get_draft_token_length +from ..scheduler import ScheduledRequests if TYPE_CHECKING: - from .resource_manager import KVCacheManager + from ..resource_manager import KVCacheManager # Used to store data for a single inflight request. @@ -86,7 +86,6 @@ class SchedulerOutput: class KvCacheConnectorWorker(ABC): - def __init__(self, llm_args: TorchLlmArgs): self._llm_args = llm_args self._metadata = None @@ -160,26 +159,31 @@ def wait_for_save(self, stream: torch.cuda.Stream): @abstractmethod def get_finished( - self, finished_gen_req_ids: List[int], - started_loading_req_ids: List[int]) -> Tuple[List[int], List[int]]: + self, finished_gen_req_ids: List[int], started_loading_req_ids: List[int] + ) -> Tuple[List[int], List[int]]: """ Get the requests that have finished loading and saving. Args: - finished_gen_req_ids: The IDs of the requests that have finished generating tokens, and are now asynchronously saving. - started_loading_req_ids: The IDs of the requests that have started asynchronously loading. + finished_gen_req_ids: The IDs of the requests that have + finished generating tokens, and are now asynchronously saving. + started_loading_req_ids: The IDs of the requests that have + started asynchronously loading. Returns: The IDs of the requests that have finished saving. The IDs of the requests that have finished loading. - Note: IDs may only be returned from this call after they've been provided in the `finished_gen_req_ids` and `started_loading_req_ids` arguments. - Additionally, the runtime will only take action based on these returned IDs once they've been returned by ALL workers. This allows some workers to take longer than others to complete the operations. + Note: IDs may only be returned from this call after they've been + provided in the ``finished_gen_req_ids`` and + ``started_loading_req_ids`` arguments. Additionally, the runtime + will only take action based on these returned IDs once they've + been returned by ALL workers. This allows some workers to take + longer than others to complete the operations. """ class KvCacheConnectorScheduler(ABC): - def __init__(self, llm_args: TorchLlmArgs): self._llm_args = llm_args super().__init__() @@ -198,8 +202,8 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput): @abstractmethod def get_num_new_matched_tokens( - self, request: LlmRequest, - num_computed_tokens: int) -> Tuple[int, bool]: + self, request: LlmRequest, num_computed_tokens: int + ) -> Tuple[int, bool]: """ Get the number of tokens that can be loaded from remote KV cache. This does not include the tokens already matched on device (indicated by `num_computed_tokens`). @@ -214,8 +218,7 @@ def get_num_new_matched_tokens( """ @abstractmethod - def request_finished(self, request: LlmRequest, - cache_block_ids: List[int]) -> bool: + def request_finished(self, request: LlmRequest, cache_block_ids: List[int]) -> bool: """ Called when a request is finished generating tokens. @@ -224,12 +227,13 @@ def request_finished(self, request: LlmRequest, Returns: Whether the request is performing asynchronous saving operations. - If true, this indicates that the kv cache manager should wait to deallocate the blocks until the saving has completed (determined by `get_finished` on the workers). + If true, this indicates that the kv cache manager should wait + to deallocate the blocks until the saving has completed + (determined by ``get_finished`` on the workers). """ @abstractmethod - def update_state_after_alloc(self, request: LlmRequest, - block_ids: List[int]): + def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): """ Called after get_num_new_matched_tokens is called to provide the block ids to the scheduler. @@ -252,7 +256,7 @@ class AsyncRequests: saving: Dict[int, LlmRequest] loading: Dict[int, LlmRequest] - def add_from(self, other: 'AsyncRequests'): + def add_from(self, other: "AsyncRequests"): """ Remove requests from the other `AsyncRequests` object, and add them to this one. """ @@ -262,8 +266,7 @@ def add_from(self, other: 'AsyncRequests'): other.saving = dict() other.loading = dict() - def extract_by_id(self, saving_ids: List[int], - loading_ids: List[int]) -> 'AsyncRequests': + def extract_by_id(self, saving_ids: List[int], loading_ids: List[int]) -> "AsyncRequests": """ Extract the requests with the given IDs from this `AsyncRequests` object. @@ -298,54 +301,61 @@ def loading_ids(self) -> Set[int]: class KvCacheConnectorSchedulerOutputRequest: - def __init__(self): self.block_ids = [] self.tokens = [] - def update_and_build_data(self, req: LlmRequest, - kv_cache_manager: "KVCacheManager"): + def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManager"): block_ids = kv_cache_manager.get_cache_indices(req) tokens = req.get_tokens(0) - new_block_ids = block_ids[len(self.block_ids):] - new_tokens = tokens[len(self.tokens):] + new_block_ids = block_ids[len(self.block_ids) :] + new_tokens = tokens[len(self.tokens) :] self.block_ids.extend(new_block_ids) self.tokens.extend(new_tokens) - if req.state in (LlmRequestState.CONTEXT_INIT, - LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS): + if req.state in ( + LlmRequestState.CONTEXT_INIT, + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + ): computed_position = req.context_current_position - num_scheduled_tokens = min(req.context_remaining_length, - req.context_chunk_size) + num_scheduled_tokens = min(req.context_remaining_length, req.context_chunk_size) else: computed_position = len(tokens) - 1 num_scheduled_tokens = 1 + get_draft_token_length( - req) # Specdec with draft tokens is not supported yet. + req + ) # Specdec with draft tokens is not supported yet. # Get retention priority for each new block only if retention config is provided # (for priority-based offload filtering) priorities = None if req.kv_cache_retention_config is not None: priorities = [ - kv_cache_manager.get_priority_by_block_id(block_id) - for block_id in new_block_ids + kv_cache_manager.get_priority_by_block_id(block_id) for block_id in new_block_ids ] - return RequestData(req.request_id, new_tokens, new_block_ids, - computed_position, num_scheduled_tokens, priorities) + return RequestData( + req.request_id, + new_tokens, + new_block_ids, + computed_position, + num_scheduled_tokens, + priorities, + ) class KvCacheConnectorSchedulerOutputManager: - def __init__(self): self.requests = defaultdict(KvCacheConnectorSchedulerOutputRequest) self.external_loads = dict() - def build_scheduler_output(self, scheduled_batch: ScheduledRequests, - new_async_requests: AsyncRequests, - kv_cache_manager: "KVCacheManager"): + def build_scheduler_output( + self, + scheduled_batch: ScheduledRequests, + new_async_requests: AsyncRequests, + kv_cache_manager: "KVCacheManager", + ): scheduler_output = SchedulerOutput() for req in scheduled_batch.context_requests: @@ -355,12 +365,12 @@ def build_scheduler_output(self, scheduled_batch: ScheduledRequests, is_new = req.request_id not in self.requests request_data = self.requests[req.request_id].update_and_build_data( - req, kv_cache_manager) + req, kv_cache_manager + ) # Don't include the connector matched tokens in the initial scheduler output. if req.request_id in self.external_loads: - request_data.computed_position -= self.external_loads[ - req.request_id] + request_data.computed_position -= self.external_loads[req.request_id] if is_new: scheduler_output.new_requests.append(request_data) @@ -369,7 +379,8 @@ def build_scheduler_output(self, scheduled_batch: ScheduledRequests, for req in scheduled_batch.generation_requests: request_data = self.requests[req.request_id].update_and_build_data( - req, kv_cache_manager) + req, kv_cache_manager + ) scheduler_output.cached_requests.append(request_data) @@ -377,8 +388,7 @@ def build_scheduler_output(self, scheduled_batch: ScheduledRequests, return scheduler_output - def record_new_matched_tokens(self, request: LlmRequest, - num_new_matched_tokens: int): + def record_new_matched_tokens(self, request: LlmRequest, num_new_matched_tokens: int): self.external_loads[request.request_id] = num_new_matched_tokens @@ -388,16 +398,19 @@ class KvCacheConnectorManager(KvCacheConnectorManagerCpp): It has the following responsibilities: 1. Managing the state of async requests (both offload and onboard) - 2. Handling MPI communication. We only run the leader on one rank, but need the results of the leader API on all ranks. + 2. Handling MPI communication. We only run the leader on one rank, + but need the results of the leader API on all ranks. Note: This class is solely an implementation detail, and is not part of the connector interface itself. When implementing a connector API, you do not need to implement this class. """ - def __init__(self, worker: KvCacheConnectorWorker, - scheduler: Optional[KvCacheConnectorScheduler]): - assert (scheduler is not None) == ( - mpi_rank() == 0), "The scheduler may only exist on rank 0!" + def __init__( + self, worker: KvCacheConnectorWorker, scheduler: Optional[KvCacheConnectorScheduler] + ): + assert (scheduler is not None) == (mpi_rank() == 0), ( + "The scheduler may only exist on rank 0!" + ) super().__init__() @@ -430,29 +443,28 @@ def _run_on_leader(self, f: Callable[[], Any]) -> Any: res = None return mpi_broadcast(res, root=0) - def get_num_new_matched_tokens(self, request: LlmRequest, - num_computed_tokens: int) -> int: + def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: int) -> int: if request.is_generation_only_request: - raise RuntimeError( - "Connector API is not supported for generation-only requests!") + raise RuntimeError("Connector API is not supported for generation-only requests!") num_tokens, load_kv_async = self._run_on_leader( - lambda: self.scheduler.get_num_new_matched_tokens( - request, num_computed_tokens)) + lambda: self.scheduler.get_num_new_matched_tokens(request, num_computed_tokens) + ) if num_tokens == 0 and load_kv_async: - raise RuntimeError( - "load_kv_async must be False when num_tokens is 0!") + raise RuntimeError("load_kv_async must be False when num_tokens is 0!") # TODO(jthomson04): This part is a bit ugly. - # When the connector indicates that a request will be loaded asynchronously, we need to suspend it's execution. - # This is problematic, since at the point when this function is called, the request has already been scheduled! - # Because of this, we need to remove it from our list of scheduled requests (see `take_scheduled_requests_pending_load`). + # When the connector indicates that a request will be loaded + # asynchronously, we need to suspend its execution. This is + # problematic, since at the point when this function is called, + # the request has already been scheduled! Because of this, we + # need to remove it from our list of scheduled requests + # (see `take_scheduled_requests_pending_load`). if load_kv_async: self.new_async_requests.loading[request.request_id] = request - self.scheduler_output_manager.record_new_matched_tokens( - request, num_tokens) + self.scheduler_output_manager.record_new_matched_tokens(request, num_tokens) request.py_num_connector_matched_tokens = num_tokens @@ -462,13 +474,14 @@ def should_add_sequence(self, request: LlmRequest) -> bool: req_id = request.request_id return req_id not in self.finished_async_loading_requests - def build_scheduler_output(self, scheduled_batch: ScheduledRequests, - kv_cache_manager: "KVCacheManager"): + def build_scheduler_output( + self, scheduled_batch: ScheduledRequests, kv_cache_manager: "KVCacheManager" + ): self._scheduler_output = self.scheduler_output_manager.build_scheduler_output( - scheduled_batch, self.new_async_requests, kv_cache_manager) + scheduled_batch, self.new_async_requests, kv_cache_manager + ) - def take_scheduled_requests_pending_load( - self, scheduled_requests: ScheduledRequests): + def take_scheduled_requests_pending_load(self, scheduled_requests: ScheduledRequests): """ Remove context requests from our list of scheduled requests that are being loaded asynchronously. This is done to prevent the runtime from attempting to load the KV cache for these requests. @@ -483,8 +496,9 @@ def take_scheduled_requests_pending_load( for key in ["context_requests_chunking", "context_requests_last_chunk"]: allowed_context_requests = [] for req in getattr(scheduled_requests, key): - # If this request is being loaded asynchronously, in addition to removing it from the list of scheduled requests, - # we also need to update it's state. + # If this request is being loaded asynchronously, in + # addition to removing it from the list of scheduled + # requests, we also need to update its state. if req.request_id in self.new_async_requests.loading.keys(): req.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS @@ -499,14 +513,14 @@ def handle_metadata(self) -> object: return metadata = self._run_on_leader( - lambda: self.scheduler.build_connector_meta(self._scheduler_output)) + lambda: self.scheduler.build_connector_meta(self._scheduler_output) + ) self._scheduler_output = None self.worker.bind_connector_meta(metadata) - def request_finished(self, req: LlmRequest, - cache_block_ids: List[int]) -> bool: + def request_finished(self, req: LlmRequest, cache_block_ids: List[int]) -> bool: """ Called when a request is finished generating tokens. @@ -514,14 +528,17 @@ def request_finished(self, req: LlmRequest, req: The request that finished generating tokens. Returns: - Whether the request is performing asynchronous saving operations. If true, we do not immediately call free_resources on the request. + Whether the request is performing asynchronous saving + operations. If true, we do not immediately call + free_resources on the request. """ if req.request_id in self.finished_async_loading_requests: del self.finished_async_loading_requests[req.request_id] saving_async = self._run_on_leader( - lambda: self.scheduler.request_finished(req, cache_block_ids)) + lambda: self.scheduler.request_finished(req, cache_block_ids) + ) # This is similar to take_scheduled_requests_pending_load. # We need to update the request's state to indicate that it's still being used, but isn't schedulable. @@ -541,21 +558,23 @@ def get_finished(self) -> List[LlmRequest]: started_loading_req_ids = list(self.new_async_requests.loading_ids) finished_gen_req_ids = list(self.new_async_requests.saving_ids) - # Add the requests to our list of outstanding (still in progress) requests. + # Add the requests to our list of outstanding (still in progress) + # requests. self.pending_async_requests.add_from(self.new_async_requests) - # Pass these newly finished requests into get_finished, and get the list of requests that have finished saving and loading. - (finished_saving, - finished_loading) = self.worker.get_finished(finished_gen_req_ids, - started_loading_req_ids) + # Pass these newly finished requests into get_finished, and get + # the list of requests that have finished saving and loading. + (finished_saving, finished_loading) = self.worker.get_finished( + finished_gen_req_ids, started_loading_req_ids + ) # Remove the requests from our pending list that have finished locally. new_local_finished_async_requests = self.pending_async_requests.extract_by_id( - finished_saving, finished_loading) + finished_saving, finished_loading + ) # Add these requests to our list of locally finished requests. - self.local_finished_async_requests.add_from( - new_local_finished_async_requests) + self.local_finished_async_requests.add_from(new_local_finished_async_requests) # Broadcast this whole list to all other workers. finished_saving = list(self.local_finished_async_requests.saving_ids) @@ -564,14 +583,13 @@ def get_finished(self) -> List[LlmRequest]: all_results = mpi_allgather((finished_saving, finished_loading)) # Find only the requests that have been reported complete by all workers. - intersect_finished_saving = set.intersection( - *[set(res[0]) for res in all_results]) - intersect_finished_loading = set.intersection( - *[set(res[1]) for res in all_results]) + intersect_finished_saving = set.intersection(*[set(res[0]) for res in all_results]) + intersect_finished_loading = set.intersection(*[set(res[1]) for res in all_results]) # Remove these requests from our list of locally finished requests. all_finished = self.local_finished_async_requests.extract_by_id( - intersect_finished_saving, intersect_finished_loading) + intersect_finished_saving, intersect_finished_loading + ) # For requests that have finished loading, move them back to the context state. for id, req in all_finished.loading.items(): @@ -590,8 +608,7 @@ def set_scheduler_output(self, scheduler_output: SchedulerOutput): self._scheduler_output = scheduler_output def layer_pre_hook(self, module, *args): - self.worker.wait_for_layer_load(module.layer_idx, - torch.cuda.current_stream()) + self.worker.wait_for_layer_load(module.layer_idx, torch.cuda.current_stream()) def layer_post_hook(self, module, *args): self.worker.save_kv_layer(module.layer_idx, torch.cuda.current_stream()) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py new file mode 100644 index 000000000000..9a00cdadd7fd --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Registry of named KV cache connector presets. + +Each entry maps a short name to the import path and class names needed +by KvCacheConnectorConfig. The connector module is NOT imported here — +it is resolved at runtime via importlib in py_executor_creator.py. +""" + +CONNECTOR_REGISTRY: dict[str, dict[str, str]] = { + "lmcache": { + "connector_module": "lmcache.integration.tensorrt_llm.tensorrt_adapter", + "connector_scheduler_class": "LMCacheKvConnectorScheduler", + "connector_worker_class": "LMCacheKvConnectorWorker", + }, + "lmcache-mp": { + "connector_module": "lmcache.integration.tensorrt_llm.tensorrt_mp_adapter", + "connector_scheduler_class": "LMCacheMPKvConnectorScheduler", + "connector_worker_class": "LMCacheMPKvConnectorWorker", + }, + "kvbm": { + "connector_module": "kvbm.trtllm_integration.connector", + "connector_scheduler_class": "DynamoKVBMConnectorLeader", + "connector_worker_class": "DynamoKVBMConnectorWorker", + }, +} diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 14f42a0ae8ac..ca53ad3c8ab0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -47,13 +47,13 @@ from ..speculative.drafter import Drafter from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.speculation_gate import SpeculationGate +from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .executor_request_queue import ExecutorRequestQueue, RequestQueueItem from .guided_decoder import GuidedDecoder from .handle_additional_outputs import HandleAdditionalOutputs from .handle_logits import HandleLogits from .hang_detector import HangDetector -from .kv_cache_connector import KvCacheConnectorManager from .kv_cache_transceiver import KvCacheTransceiver from .llm_request import (ExecutorRequest, LlmRequest, LlmRequestState, LlmResponse, get_draft_token_length) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index e0aa739d8697..9c8aede18d0b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -37,9 +37,9 @@ create_py_executor_instance, instantiate_sampler, is_mla, validate_feature_combination) from .config_utils import is_nemotron_hybrid, is_qwen3_hybrid +from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import CapturableGuidedDecoder, GuidedDecoder -from .kv_cache_connector import KvCacheConnectorManager from .model_engine import PyTorchModelEngine from .model_loader import ModelLoader, _construct_checkpoint_loader from .py_executor import PyExecutor diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 971ff8c5402d..2813e7d96bac 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -46,7 +46,7 @@ from ..._utils import binding_to_str_dtype, mpi_rank, nvtx_range from ...logger import logger from ...mapping import CpType, Mapping -from .kv_cache_connector import KvCacheConnectorManager +from .connectors.kv_cache_connector import KvCacheConnectorManager from .llm_request import (LlmRequest, LlmRequestState, SamplingConfig, get_draft_token_length) from .scheduler import ScheduledRequests diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 9f001b4e5aeb..8b7b06dcba41 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -948,16 +948,55 @@ def num_capture_layers(self) -> int: class KvCacheConnectorConfig(StrictBaseModel): """ Configuration for the KV Cache Connector. + + Can be configured either by specifying a named preset via ``connector`` + (e.g. ``"lmcache"``), or by providing explicit ``connector_module``, + ``connector_scheduler_class``, and ``connector_worker_class`` fields. + When ``connector`` is set, the module/class fields are auto-populated + from the preset registry and can be omitted. """ - connector_module: str = Field( - ..., + connector: Optional[str] = Field( + None, + description="Named connector preset (e.g. 'lmcache'). " + "When set, connector_module/scheduler_class/worker_class are " + "auto-populated from the preset registry.") + connector_module: Optional[str] = Field( + None, description= "The import path to the connector module. It will be imported with `importlib.import_module`." ) - connector_scheduler_class: str = Field( - ..., description="The class name of the scheduler within the module.") - connector_worker_class: str = Field( - ..., description="The class name of the worker within the module.") + connector_scheduler_class: Optional[str] = Field( + None, description="The class name of the scheduler within the module.") + connector_worker_class: Optional[str] = Field( + None, description="The class name of the worker within the module.") + server_url: Optional[str] = Field( + None, + description="URL for an external connector server " + "(e.g. 'tcp://localhost:5555'). Connectors that run in " + "multi-process mode use this to reach the cache server.") + + @model_validator(mode="after") + def _resolve_preset(self) -> "KvCacheConnectorConfig": + from tensorrt_llm._torch.pyexecutor.connectors.registry import \ + CONNECTOR_REGISTRY + if self.connector is not None: + preset = CONNECTOR_REGISTRY.get(self.connector) + if preset is None: + raise ValueError( + f"Unknown connector preset: {self.connector!r}. " + f"Known presets: {list(CONNECTOR_REGISTRY)}") + for k, v in preset.items(): + if getattr(self, k) is None: + object.__setattr__(self, k, v) + if self.connector_module is None: + raise ValueError( + "connector_module is required (set 'connector' to use a " + "named preset, or provide connector_module explicitly)") + if self.connector_scheduler_class is None: + raise ValueError("connector_scheduler_class is required") + if self.connector_worker_class is None: + raise ValueError("connector_worker_class is required") + return self class LayerwiseBenchmarksConfig(StrictBaseModel): diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 494c803ac7e2..c683c7f76694 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -6,7 +6,7 @@ import tensorrt_llm from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.pyexecutor.kv_cache_connector import \ +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ KvCacheConnectorWorker from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine diff --git a/tests/unittest/_torch/test_connector.py b/tests/unittest/_torch/test_connector.py index 5a70ac639557..96cdbe9dc393 100644 --- a/tests/unittest/_torch/test_connector.py +++ b/tests/unittest/_torch/test_connector.py @@ -22,7 +22,7 @@ import pytest from tensorrt_llm import mpi_rank -from tensorrt_llm._torch.pyexecutor.kv_cache_connector import ( +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( AsyncRequests, KvCacheConnectorManager, KvCacheConnectorSchedulerOutputManager) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState