LMCache 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"""Demonstrates using LMCache as a KV cache backend for TensorRT-LLM.
 16
 17Uses the KV Cache Connector interface.
 18
 19LMCache stores previously computed KV tensors and replays them on subsequent
 20requests with the same prefix, reducing recomputation.
 21
 22The connector implementation lives in LMCache:
 23  lmcache.integration.tensorrt_llm.tensorrt_adapter
 24
 25TRT-LLM resolves the ``"lmcache"`` preset to the correct import paths
 26automatically via the connector registry.
 27
 28Prerequisites:
 29  pip install lmcache
 30
 31How to run:
 32  PYTHONHASHSEED=0 LMCACHE_CHUNK_SIZE=32 \
 33    python llm_lmcache_connector.py Qwen/Qwen2-1.5B-Instruct
 34
 35Note: PYTHONHASHSEED=0 must be set before the Python process starts
 36to ensure deterministic cache key hashing in LMCache.
 37
 38Expected output:
 39  The first request logs "Stored N ... tokens". The second request logs
 40  "Retrieved N ... tokens" with N > 0, then the script prints
 41  "OK: outputs match." The LMCache retrieval log is the cache-hit signal.
 42
 43See Also:
 44  examples/llm-api/configs/trtllm_lmcache_connector_extra.yaml -- trtllm-serve YAML
 45"""
 46
 47import click
 48
 49from tensorrt_llm import LLM, SamplingParams
 50from tensorrt_llm.llmapi.llm_args import KvCacheConfig, KvCacheConnectorConfig
 51
 52try:
 53    from lmcache.integration.tensorrt_llm import destroy_engine
 54except ImportError as e:
 55    raise ImportError(
 56        "LMCache is not installed or is missing the TensorRT-LLM integration. "
 57        "Run: pip install 'lmcache'"
 58    ) from e
 59
 60# A prompt long enough to produce at least one full TRT-LLM KV block.
 61_TEST_PROMPT = (
 62    "Nvidia Corporation is an American technology company headquartered in "
 63    "Santa Clara, California. Founded in 1993 by Jensen Huang, Chris "
 64    "Malachowsky, and Curtis Priem, it develops graphics processing units "
 65    "(GPUs), system on a chips (SoCs), and application programming "
 66    "interfaces (APIs) for data science, high-performance computing, and "
 67    "mobile and automotive applications. Tell me about the company."
 68)
 69
 70
 71@click.command()
 72@click.argument("model", type=str)
 73def main(model: str):
 74    # Match LMCACHE_CHUNK_SIZE and disable TensorRT-LLM's in-GPU block reuse.
 75    # Otherwise it can serve the repeated prompt without exercising LMCache.
 76    kv_cache_config = KvCacheConfig(enable_block_reuse=False, tokens_per_block=32)
 77    kv_connector_config = KvCacheConnectorConfig(connector="lmcache")
 78    sampling_params = SamplingParams(max_tokens=32)
 79
 80    # The in-process LMCache engine is scoped to the LLM instance, so both
 81    # requests use the same instance.
 82    llm = LLM(
 83        model=model,
 84        backend="pytorch",
 85        kv_cache_config=kv_cache_config,
 86        kv_connector_config=kv_connector_config,
 87    )
 88
 89    print("--- First request (cold LMCache; KV will be stored) ---")
 90    output0 = llm.generate([_TEST_PROMPT], sampling_params)
 91    text0 = output0[0].outputs[0].text
 92    print("First output:", text0)
 93
 94    print("\n--- Second request (warm LMCache; watch for 'Retrieved N ... tokens') ---")
 95    output1 = llm.generate([_TEST_PROMPT], sampling_params)
 96    text1 = output1[0].outputs[0].text
 97    print("Second output:", text1)
 98
 99    assert text0 == text1, (
100        f"Outputs differ between identical requests.\nFirst:  {text0!r}\nSecond: {text1!r}"
101    )
102    print("\nOK: outputs match. The LMCache retrieval log above confirms external KV reuse.")
103
104    destroy_engine()
105
106
107if __name__ == "__main__":
108    main()