KV Cache Connector#

Source NVIDIA/TensorRT-LLM.

  1# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  2# SPDX-License-Identifier: Apache-2.0
  3#
  4# Licensed under the Apache License, Version 2.0 (the "License");
  5# you may not use this file except in compliance with the License.
  6# You may obtain a copy of the License at
  7#
  8# http://www.apache.org/licenses/LICENSE-2.0
  9#
 10# Unless required by applicable law or agreed to in writing, software
 11# distributed under the License is distributed on an "AS IS" BASIS,
 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13# See the License for the specific language governing permissions and
 14# limitations under the License.
 15
 16'''
 17This script demonstrates the KV cache connector feature in TensorRT-LLM, which enables
 18custom persistence and reuse of KV cache blocks across different LLM instances.
 19
 20**Scenario:**
 21The script implements a persistent KV cache connector that saves computed KV cache blocks
 22to disk and loads them back in subsequent runs, eliminating redundant computation for
 23recurring prompts.
 24
 25**What is a KV Cache Connector?**
 26
 27A KV cache connector is a customizable interface that allows you to:
 281.  **Save KV Cache:** Persist computed KV cache blocks to an external storage
 29    (disk, database, distributed cache, etc.)
 302.  **Load KV Cache:** Retrieve previously computed cache blocks instead of recomputing them
 313.  **Share Cache Across Instances:** Reuse cache blocks across different LLM instances
 32    or sessions, unlike regular block reuse which is limited to a single instance
 33
 34**How It Works:**
 35
 36This example implements a `PersistentKvCacheConnector` with two key components:
 37
 38* **PersistentKvCacheConnectorLeader (Scheduler):**
 39    - Hashes token sequences to create unique identifiers for each cache block
 40    - Checks if cached blocks exist on disk for incoming requests
 41    - Schedules load operations for cache hits
 42    - Schedules save operations for newly computed blocks
 43
 44* **PersistentKvCacheConnectorWorker:**
 45    - Executes the actual load/save operations between GPU and disk
 46    - Loads cached blocks from disk files into GPU memory
 47    - Saves newly computed blocks from GPU to disk files
 48
 49**Demonstration:**
 50
 51The script processes the same prompt twice using two separate LLM instances:
 52
 531.  **First Run (Instance 1):**
 54    - The LLM computes the KV cache for the input prompt
 55    - The connector saves the computed cache blocks to disk (as .pt files)
 56    - The generation completes and the LLM instance is destroyed
 57
 582.  **Second Run (Instance 2):**
 59    - A new LLM instance is created with the same connector configuration
 60    - When processing the same prompt, the connector finds matching cache blocks on disk
 61    - The cache is loaded from disk instead of being recomputed
 62    - **Expected Outcome:** Faster prefill as cache blocks are loaded rather than computed
 63    - Both outputs should be identical, demonstrating deterministic cache reuse
 64
 65**Key Benefits:**
 66
 67- **Cross-Instance Cache Sharing:** Share computed caches across multiple LLM instances
 68- **Persistent Storage:** Cache survives beyond the lifetime of a single LLM instance
 69- **Custom Storage Backends:** Implement any storage mechanism (shown here: disk files)
 70- **Reduced Computation:** Eliminate redundant KV cache computation for repeated prompts
 71
 72**How to Run:**
 73
 74```bash
 75python llm_kv_cache_connector.py <model_path>
 76```
 77
 78Example:
 79```bash
 80python llm_kv_cache_connector.py meta-llama/Llama-3.1-8B-Instruct
 81```
 82
 83**Implementation Notes:**
 84
 85- This example uses content-based hashing to identify cache blocks
 86- Cache files are stored in a temporary directory (cleaned up after the demo)
 87- The implementation is simplified and not optimized for production use
 88- Does not support chunked prefill in this example
 89- See `tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py` for the full connector interface
 90
 91**NOTE:** This example connector implementation is designed for demonstration purposes
 92and is NOT suitable for production use without additional optimizations and error handling.
 93'''
 94
 95import hashlib
 96import os
 97import sys
 98from dataclasses import dataclass, field
 99from pathlib import Path
100from tempfile import NamedTemporaryFile, TemporaryDirectory
101from typing import Optional
102
103import click
104import torch
105
106from tensorrt_llm import LLM, SamplingParams, logger
107from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import (
108    KvCacheConnectorScheduler, KvCacheConnectorWorker, SchedulerOutput)
109from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import \
110    valid_page_slots
111from tensorrt_llm.bindings.internal.batch_manager import LlmRequest
112from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig, TorchLlmArgs
113
114CONNECTOR_CACHE_FOLDER_KEY = "CONNECTOR_CACHE_FOLDER"
115
116
117@dataclass
118class PersistentKvCacheConnectorMetadata:
119    load: list[tuple[str, int]] = field(default_factory=list)
120    save: list[tuple[str, int]] = field(default_factory=list)
121    prefix_load_ids: list[int] = field(default_factory=list)
122
123
124class PersistentKvCacheConnectorWorker(KvCacheConnectorWorker):
125
126    def __init__(self, llm_args: TorchLlmArgs):
127        super().__init__(llm_args)
128
129        self.kv_cache_tensor = None
130        self._finished_prefix_loads: list[int] = []
131
132    def register_kv_caches(self, kv_cache_tensor: torch.Tensor):
133        # This is the only registration hook this connector needs. A cache that
134        # describes itself as a layout instead still arrives here, through
135        # `register_kv_cache_layout`'s default, as long as one tensor can
136        # describe it. See llm_kv_cache_connector_vswa.py for the case where it
137        # cannot -- one attention window size per layer group.
138        assert self.kv_cache_tensor is None, "KV cache tensor already registered"
139        self.kv_cache_tensor = kv_cache_tensor
140
141    def start_load_kv(self, stream: torch.cuda.Stream):
142        # Do all loads synchronously, and blockwise.
143        for path, block_id in self._metadata.load:
144            cpu_tensor = torch.load(path, map_location="cpu")
145
146            # Copy into the device block.
147            self.kv_cache_tensor[block_id].copy_(cpu_tensor, non_blocking=False)
148
149        self._finished_prefix_loads.extend(self._metadata.prefix_load_ids)
150
151    def get_finished_prefix_loads(self) -> list[int]:
152        finished, self._finished_prefix_loads = self._finished_prefix_loads, []
153        return finished
154
155    def wait_for_layer_load(self, layer_idx: int, stream: torch.cuda.Stream):
156        pass
157
158    def save_kv_layer(self, layer_idx: int, stream: torch.cuda.Stream):
159        pass
160
161    def wait_for_save(self, stream: torch.cuda.Stream):
162
163        # Make sure the forward pass is complete before beginning our save.
164        stream.synchronize()
165
166        for path, block_id in self._metadata.save:
167            cpu_tensor = self.kv_cache_tensor[block_id].cpu()
168
169            # Don't write anything if this specific block already exists.
170            if Path(path).exists():
171                continue
172
173            # Publish complete immutable files so a reservation can retain an
174            # inode while another process replaces or removes its cache key.
175            with NamedTemporaryFile(dir=Path(path).parent) as staging:
176                torch.save(cpu_tensor, staging.name)
177                try:
178                    os.link(staging.name, path)
179                except FileExistsError:
180                    pass
181
182    def get_finished(
183            self, finished_gen_req_ids: list[int],
184            started_loading_req_ids: list[int]) -> tuple[list[int], list[int]]:
185
186        return [], []
187
188
189class PersistentKvCacheConnectorLeader(KvCacheConnectorScheduler):
190
191    def __init__(self, llm_args: TorchLlmArgs):
192        super().__init__(llm_args)
193
194        self.block_size = self._llm_args.kv_cache_config.tokens_per_block
195        self.pending_loads = {}
196
197        self.cache_folder = os.environ.get(CONNECTOR_CACHE_FOLDER_KEY,
198                                           "./connector_cache")
199
200        os.makedirs(self.cache_folder, exist_ok=True)
201        self._reservation_folder = TemporaryDirectory(prefix=".reservations-",
202                                                      dir=self.cache_folder)
203        self._reserved_files: dict[int, dict[int, Path]] = {}
204        self._reserved_ranges: dict[int, list[tuple[int, int]]] = {}
205
206    def build_connector_meta(self, scheduler_output: SchedulerOutput):
207        # NOTE: This is a simplified implementation, and does not work with chunked prefill.
208
209        metadata = PersistentKvCacheConnectorMetadata()
210
211        accepted_ends = {}
212        for load in scheduler_output.prefix_loads:
213            if len(load.block_ids_by_layer_group) != 1:
214                raise ValueError(
215                    "Persistent connector requires one layer group")
216            if load.start % self.block_size or load.end % self.block_size:
217                raise ValueError("Persistent connector loads complete blocks")
218            block_ids = load.block_ids_by_layer_group[0]
219            if load.end // self.block_size > len(block_ids):
220                raise ValueError(
221                    "Confirmed prefix exceeds destination allocation")
222            slots = dict(valid_page_slots(block_ids))
223            files = self._reserved_files[load.reservation_id]
224            for ordinal in range(load.start // self.block_size,
225                                 load.end // self.block_size):
226                if ordinal in slots:
227                    metadata.load.append((str(files[ordinal]), slots[ordinal]))
228            metadata.prefix_load_ids.append(load.reservation_id)
229            accepted_ends[load.request_id] = load.end
230
231        for req in scheduler_output.new_requests:
232            num_computed_blocks = req.computed_position // self.block_size
233            slots = dict(valid_page_slots(req.new_block_ids))
234            pending_load = self.pending_loads.get(req.request_id, [])
235            for ordinal, path in enumerate(pending_load, num_computed_blocks):
236                if ordinal in slots:
237                    metadata.load.append((str(path), slots[ordinal]))
238
239            loaded_end = accepted_ends.get(
240                req.request_id,
241                (num_computed_blocks + len(pending_load)) * self.block_size)
242            computed_end = max(req.computed_position,
243                               loaded_end) + req.num_scheduled_tokens
244            for ordinal, slot in slots.items():
245                end = (ordinal + 1) * self.block_size
246                if end <= loaded_end or end > min(len(req.new_tokens),
247                                                  computed_end):
248                    continue
249                key = self._hash_tokens(req.new_tokens[:end], req.cache_salt)
250                metadata.save.append((str(self._file_path(key)), slot))
251
252        self.pending_loads = {}
253        return metadata
254
255    def _hash_tokens(self, tokens: list[int], cache_salt: Optional[str]) -> str:
256        # KV depends on every preceding token, including those in other blocks.
257        key = repr((cache_salt, tuple(tokens))).encode("utf-8")
258        return hashlib.sha256(key).hexdigest()
259
260    def _file_path(self, hash_value: str) -> Path:
261        return Path(self.cache_folder) / f"{hash_value}.pt"
262
263    def get_num_new_matched_tokens(
264            self, request: LlmRequest,
265            num_computed_tokens: int) -> tuple[int, bool]:
266        self.pending_loads[request.request_id] = []
267
268        # Don't bother with sequences with partial matches.
269        if (num_computed_tokens % self.block_size) != 0:
270            return 0, False
271
272        tokens = request.get_tokens(0)
273        for end in range(num_computed_tokens + self.block_size,
274                         len(tokens) + 1, self.block_size):
275            key = self._hash_tokens(tokens[:end], request.cache_salt)
276            file_path = self._file_path(key)
277            if not file_path.exists():
278                break
279            self.pending_loads[request.request_id].append(file_path)
280
281        logger.info(
282            f"KV CONNECTOR: Matched {len(self.pending_loads[request.request_id])} blocks for request {request.request_id}"
283        )
284
285        return len(
286            self.pending_loads[request.request_id]) * self.block_size, False
287
288    def reserve_prefix(self, request: LlmRequest, num_computed_tokens: int,
289                       reservation_id: int) -> tuple[int, bool]:
290        """Protect immutable source files without reading their KV data."""
291        if num_computed_tokens % self.block_size:
292            return 0, False
293        tokens = request.get_tokens(0)
294        folder = Path(self._reservation_folder.name) / str(reservation_id)
295        folder.mkdir()
296        files = {}
297        for end in range(num_computed_tokens + self.block_size,
298                         len(tokens) + 1, self.block_size):
299            ordinal = end // self.block_size - 1
300            key = self._hash_tokens(tokens[:end], request.cache_salt)
301            protected = folder / f"{ordinal}.pt"
302            try:
303                os.link(self._file_path(key), protected)
304            except FileNotFoundError:
305                break
306            files[ordinal] = protected
307        count = len(files) * self.block_size
308        if count:
309            self._reserved_files[reservation_id] = files
310            self._reserved_ranges[reservation_id] = [
311                (num_computed_tokens, num_computed_tokens + count)
312            ]
313        else:
314            folder.rmdir()
315        return count, False
316
317    def release_prefix_reservation(self, request: LlmRequest,
318                                   reservation_id: int, start: int,
319                                   end: int) -> None:
320        """Release only the named range; overlapping reservations retain it."""
321        remaining = []
322        for left, right in self._reserved_ranges[reservation_id]:
323            if right <= start or left >= end:
324                remaining.append((left, right))
325            else:
326                if left < start:
327                    remaining.append((left, start))
328                if end < right:
329                    remaining.append((end, right))
330        files = self._reserved_files[reservation_id]
331        for ordinal, path in list(files.items()):
332            block_start = ordinal * self.block_size
333            block_end = block_start + self.block_size
334            if not any(left < block_end and right > block_start
335                       for left, right in remaining):
336                path.unlink()
337                del files[ordinal]
338        if remaining:
339            self._reserved_ranges[reservation_id] = remaining
340        else:
341            del self._reserved_ranges[reservation_id]
342            del self._reserved_files[reservation_id]
343            (Path(self._reservation_folder.name) / str(reservation_id)).rmdir()
344
345    def request_finished(self, request: LlmRequest,
346                         cache_block_ids: list[int]) -> bool:
347        # We don't do any asynchronous saving, so always return False
348        return False
349
350    def update_state_after_alloc(self, request: LlmRequest,
351                                 block_ids: list[int]):
352        pass
353
354
355@click.command()
356@click.argument("model", type=str)
357def main(model: str):
358    sys.path.append(os.path.join(
359        os.path.dirname(__file__),
360        "..",
361    ))
362
363    this_module = __file__[__file__.rfind("/") + 1:__file__.rfind(".py")]
364
365    # --- KV Cache Connector Config ---
366    kv_connector_config = KvCacheConnectorConfig(
367        connector_module=this_module,
368        connector_scheduler_class="PersistentKvCacheConnectorLeader",
369        connector_worker_class="PersistentKvCacheConnectorWorker",
370    )
371
372    connector_cache_dir = TemporaryDirectory()
373    os.environ[CONNECTOR_CACHE_FOLDER_KEY] = connector_cache_dir.name
374
375    # Create LLM instance with KV Cache Connector
376    llm = LLM(model=model,
377              backend="pytorch",
378              cuda_graph_config=None,
379              kv_connector_config=kv_connector_config)
380
381    test_text = (
382        "Nvidia Corporation is an American technology company headquartered in Santa Clara, California."
383        "Founded in 1993 by Jensen Huang, Chris Malachowsky, and Curtis Priem, it develops graphics processing units (GPUs), "
384        "system on a chips (SoCs), and application programming interfaces (APIs) for data science, high-performance computing, "
385        "and mobile and automotive applications. Tell me about the company.")
386
387    sampling_params = SamplingParams(max_tokens=32)
388
389    # Generate text with the first LLM instance and save the kv cache blocks by the connector.
390    output = llm.generate([test_text], sampling_params)
391    text0 = output[0].outputs[0].text
392
393    print("First output: ", text0)
394    print("Loading new LLM instance...")
395
396    del llm
397
398    # Create a new LLM instance with the same connector configuration
399    llm = LLM(model=model,
400              backend="pytorch",
401              cuda_graph_config=None,
402              kv_connector_config=kv_connector_config)
403
404    # Generate text with the second LLM instance and it should reuse the kv cache blocks from the connector.
405    output = llm.generate([test_text], sampling_params)
406    text1 = output[0].outputs[0].text
407
408    print("Second output (using connector cache): ", text1)
409
410    # Verify that the two outputs are identical
411    assert text0 == text1
412
413    connector_cache_dir.cleanup()
414
415
416if __name__ == "__main__":
417    main()