KV Cache Connector#

Source NVIDIA/TensorRT-LLM.

  1'''
  2This script demonstrates the KV cache connector feature in TensorRT-LLM, which enables
  3custom persistence and reuse of KV cache blocks across different LLM instances.
  4
  5**Scenario:**
  6The script implements a persistent KV cache connector that saves computed KV cache blocks
  7to disk and loads them back in subsequent runs, eliminating redundant computation for
  8recurring prompts.
  9
 10**What is a KV Cache Connector?**
 11
 12A KV cache connector is a customizable interface that allows you to:
 131.  **Save KV Cache:** Persist computed KV cache blocks to an external storage
 14    (disk, database, distributed cache, etc.)
 152.  **Load KV Cache:** Retrieve previously computed cache blocks instead of recomputing them
 163.  **Share Cache Across Instances:** Reuse cache blocks across different LLM instances
 17    or sessions, unlike regular block reuse which is limited to a single instance
 18
 19**How It Works:**
 20
 21This example implements a `PersistentKvCacheConnector` with two key components:
 22
 23* **PersistentKvCacheConnectorLeader (Scheduler):**
 24    - Hashes token sequences to create unique identifiers for each cache block
 25    - Checks if cached blocks exist on disk for incoming requests
 26    - Schedules load operations for cache hits
 27    - Schedules save operations for newly computed blocks
 28
 29* **PersistentKvCacheConnectorWorker:**
 30    - Executes the actual load/save operations between GPU and disk
 31    - Loads cached blocks from disk files into GPU memory
 32    - Saves newly computed blocks from GPU to disk files
 33
 34**Demonstration:**
 35
 36The script processes the same prompt twice using two separate LLM instances:
 37
 381.  **First Run (Instance 1):**
 39    - The LLM computes the KV cache for the input prompt
 40    - The connector saves the computed cache blocks to disk (as .pt files)
 41    - The generation completes and the LLM instance is destroyed
 42
 432.  **Second Run (Instance 2):**
 44    - A new LLM instance is created with the same connector configuration
 45    - When processing the same prompt, the connector finds matching cache blocks on disk
 46    - The cache is loaded from disk instead of being recomputed
 47    - **Expected Outcome:** Faster prefill as cache blocks are loaded rather than computed
 48    - Both outputs should be identical, demonstrating deterministic cache reuse
 49
 50**Key Benefits:**
 51
 52- **Cross-Instance Cache Sharing:** Share computed caches across multiple LLM instances
 53- **Persistent Storage:** Cache survives beyond the lifetime of a single LLM instance
 54- **Custom Storage Backends:** Implement any storage mechanism (shown here: disk files)
 55- **Reduced Computation:** Eliminate redundant KV cache computation for repeated prompts
 56
 57**How to Run:**
 58
 59```bash
 60python llm_kv_cache_connector.py <model_path>
 61```
 62
 63Example:
 64```bash
 65python llm_kv_cache_connector.py meta-llama/Llama-3.1-8B-Instruct
 66```
 67
 68**Implementation Notes:**
 69
 70- This example uses content-based hashing to identify cache blocks
 71- Cache files are stored in a temporary directory (cleaned up after the demo)
 72- The implementation is simplified and not optimized for production use
 73- Does not support chunked prefill in this example
 74- See `tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py` for the full connector interface
 75
 76**NOTE:** This example connector implementation is designed for demonstration purposes
 77and is NOT suitable for production use without additional optimizations and error handling.
 78'''
 79
 80import os
 81import sys
 82from dataclasses import dataclass, field
 83from pathlib import Path
 84from tempfile import TemporaryDirectory
 85from typing import Optional
 86
 87import click
 88import torch
 89
 90from tensorrt_llm import LLM, SamplingParams, logger
 91from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import (
 92    KvCacheConnectorScheduler, KvCacheConnectorWorker, SchedulerOutput)
 93from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import \
 94    valid_page_slots
 95from tensorrt_llm.bindings.internal.batch_manager import LlmRequest
 96from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig, TorchLlmArgs
 97
 98CONNECTOR_CACHE_FOLDER_KEY = "CONNECTOR_CACHE_FOLDER"
 99
100
101@dataclass
102class PersistentKvCacheConnectorMetadata:
103    load: list[tuple[str, int]] = field(default_factory=list)
104    save: list[tuple[str, int]] = field(default_factory=list)
105
106
107class PersistentKvCacheConnectorWorker(KvCacheConnectorWorker):
108
109    def __init__(self, llm_args: TorchLlmArgs):
110        super().__init__(llm_args)
111
112        self.kv_cache_tensor = None
113
114    def register_kv_caches(self, kv_cache_tensor: torch.Tensor):
115        # This is the only registration hook this connector needs. A cache that
116        # describes itself as a layout instead still arrives here, through
117        # `register_kv_cache_layout`'s default, as long as one tensor can
118        # describe it. See llm_kv_cache_connector_vswa.py for the case where it
119        # cannot -- one attention window size per layer group.
120        assert self.kv_cache_tensor is None, "KV cache tensor already registered"
121        self.kv_cache_tensor = kv_cache_tensor
122
123    def start_load_kv(self, stream: torch.cuda.Stream):
124        # Do all loads synchronously, and blockwise.
125        for path, block_id in self._metadata.load:
126            cpu_tensor = torch.load(path, map_location="cpu")
127
128            # Copy into the device block.
129            self.kv_cache_tensor[block_id].copy_(cpu_tensor, non_blocking=False)
130
131    def wait_for_layer_load(self, layer_idx: int, stream: torch.cuda.Stream):
132        pass
133
134    def save_kv_layer(self, layer_idx: int, stream: torch.cuda.Stream):
135        pass
136
137    def wait_for_save(self, stream: torch.cuda.Stream):
138
139        # Make sure the forward pass is complete before beginning our save.
140        stream.synchronize()
141
142        for path, block_id in self._metadata.save:
143            cpu_tensor = self.kv_cache_tensor[block_id].cpu()
144
145            # Don't write anything if this specific block already exists.
146            if Path(path).exists():
147                continue
148
149            # Do a blocking save to the file. This way, we only return once all saves are complete.
150            torch.save(cpu_tensor, path)
151
152    def get_finished(
153            self, finished_gen_req_ids: list[int],
154            started_loading_req_ids: list[int]) -> tuple[list[int], list[int]]:
155
156        return [], []
157
158
159class PersistentKvCacheConnectorLeader(KvCacheConnectorScheduler):
160
161    def __init__(self, llm_args: TorchLlmArgs):
162        super().__init__(llm_args)
163
164        self.block_size = self._llm_args.kv_cache_config.tokens_per_block
165        self.pending_loads = {}
166
167        self.cache_folder = os.environ.get(CONNECTOR_CACHE_FOLDER_KEY,
168                                           "./connector_cache")
169
170        os.makedirs(self.cache_folder, exist_ok=True)
171
172    def build_connector_meta(self, scheduler_output: SchedulerOutput):
173        # NOTE: This is a simplified implementation, and does not work with chunked prefill.
174
175        metadata = PersistentKvCacheConnectorMetadata()
176
177        for req in scheduler_output.new_requests:
178            # If we don't have any pending loads for this request, we can skip it.
179            if req.request_id not in self.pending_loads:
180                continue
181
182            num_computed_blocks = req.computed_position // self.block_size
183            block_ids = req.new_block_ids
184
185            pending_load = self.pending_loads[req.request_id]
186
187            # Ordinal -> page slot for the blocks that have a page. Blocks with
188            # none keep their ordinal in `block_ids` so that entry `i` always
189            # describes the same token range; they are dropped here so no
190            # transfer can be built against one.
191            slots = dict(valid_page_slots(block_ids))
192
193            for file_path, block_pos in zip(
194                    pending_load, range(num_computed_blocks, len(block_ids))):
195                slot = slots.get(block_pos)
196                if slot is None:
197                    continue
198                metadata.load.append((file_path, slot))
199
200            # Break up the remainder of the token sequence into chunks.
201            chunks = self._chunk_tokens(req.new_tokens)
202
203            # For each chunk that isn't already on device, and isn't in our connector cache, we need to save it.
204            for block_pos in range(num_computed_blocks + len(pending_load),
205                                   len(block_ids)):
206                slot = slots.get(block_pos)
207                if slot is None:
208                    continue
209                if len(chunks[block_pos]) == self.block_size:
210                    hashed_tokens = self._hash_tokens(chunks[block_pos],
211                                                      req.cache_salt)
212
213                    file_path = self._file_path(hashed_tokens)
214
215                    metadata.save.append((file_path, slot))
216
217        self.pending_loads = {}
218
219        return metadata
220
221    def _hash_tokens(self, tokens: list[int], cache_salt: Optional[str]) -> int:
222        # cache_salt must participate in the hash so that requests carrying
223        # different salts (or no salt) cannot collide on the same cache file.
224        return abs(hash((cache_salt, tuple(tokens))))
225
226    def _file_path(self, hash_value: int) -> Path:
227        return Path(self.cache_folder) / f"{hash_value}.pt"
228
229    def _chunk_tokens(self, tokens: list[int]) -> list[list[int]]:
230        return [
231            tokens[i:i + self.block_size]
232            for i in range(0, len(tokens), self.block_size)
233        ]
234
235    def get_num_new_matched_tokens(
236            self, request: LlmRequest,
237            num_computed_tokens: int) -> tuple[int, bool]:
238        self.pending_loads[request.request_id] = []
239
240        # Don't bother with sequences with partial matches.
241        if (num_computed_tokens % self.block_size) != 0:
242            return 0, False
243
244        computed_blocks = num_computed_tokens // self.block_size
245
246        # Get all the tokens that don't have a cache hit on device.
247        remaining_tokens = request.get_tokens(0)[computed_blocks *
248                                                 self.block_size:]
249
250        remaining_chunks = self._chunk_tokens(remaining_tokens)
251
252        # For each chunk, check if it exists in our cache.
253        for chunk in remaining_chunks:
254            # Only do full blocks.
255            if len(chunk) == self.block_size:
256                hashed_tokens = self._hash_tokens(chunk, request.cache_salt)
257
258                file_path = self._file_path(hashed_tokens)
259
260                # If we get a cache hit, we want to load it into device.
261                # Otherwise, we can stop looking.
262                if file_path.exists():
263                    self.pending_loads[request.request_id].append(file_path)
264                else:
265                    break
266
267        logger.info(
268            f"KV CONNECTOR: Matched {len(self.pending_loads[request.request_id])} blocks for request {request.request_id}"
269        )
270
271        return len(
272            self.pending_loads[request.request_id]) * self.block_size, False
273
274    def request_finished(self, request: LlmRequest,
275                         cache_block_ids: list[int]) -> bool:
276        # We don't do any asynchronous saving, so always return False
277        return False
278
279    def update_state_after_alloc(self, request: LlmRequest,
280                                 block_ids: list[int]):
281        pass
282
283
284@click.command()
285@click.argument("model", type=str)
286def main(model: str):
287    sys.path.append(os.path.join(
288        os.path.dirname(__file__),
289        "..",
290    ))
291
292    this_module = __file__[__file__.rfind("/") + 1:__file__.rfind(".py")]
293
294    # --- KV Cache Connector Config ---
295    kv_connector_config = KvCacheConnectorConfig(
296        connector_module=this_module,
297        connector_scheduler_class="PersistentKvCacheConnectorLeader",
298        connector_worker_class="PersistentKvCacheConnectorWorker",
299    )
300
301    connector_cache_dir = TemporaryDirectory()
302    os.environ[CONNECTOR_CACHE_FOLDER_KEY] = connector_cache_dir.name
303
304    # Create LLM instance with KV Cache Connector
305    llm = LLM(model=model,
306              backend="pytorch",
307              cuda_graph_config=None,
308              kv_connector_config=kv_connector_config)
309
310    test_text = (
311        "Nvidia Corporation is an American technology company headquartered in Santa Clara, California."
312        "Founded in 1993 by Jensen Huang, Chris Malachowsky, and Curtis Priem, it develops graphics processing units (GPUs), "
313        "system on a chips (SoCs), and application programming interfaces (APIs) for data science, high-performance computing, "
314        "and mobile and automotive applications. Tell me about the company.")
315
316    sampling_params = SamplingParams(max_tokens=32)
317
318    # Generate text with the first LLM instance and save the kv cache blocks by the connector.
319    output = llm.generate([test_text], sampling_params)
320    text0 = output[0].outputs[0].text
321
322    print("First output: ", text0)
323    print("Loading new LLM instance...")
324
325    del llm
326
327    # Create a new LLM instance with the same connector configuration
328    llm = LLM(model=model,
329              backend="pytorch",
330              cuda_graph_config=None,
331              kv_connector_config=kv_connector_config)
332
333    # Generate text with the second LLM instance and it should reuse the kv cache blocks from the connector.
334    output = llm.generate([test_text], sampling_params)
335    text1 = output[0].outputs[0].text
336
337    print("Second output (using connector cache): ", text1)
338
339    # Verify that the two outputs are identical
340    assert text0 == text1
341
342    connector_cache_dir.cleanup()
343
344
345if __name__ == "__main__":
346    main()