KV Cache Compression#

Source NVIDIA/TensorRT-LLM.

  1# SPDX-FileCopyrightText: Copyright (c) 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.
 15r"""Configure KV cache compression with TensorRT-LLM.
 16
 17This example shows the two available compression methods. Currently, only one
 18KV cache compression method can be enabled for each LLM instance.
 19
 20NVFP4 cold-page quantization
 21----------------------------
 22Attention KV keeps its runtime dtype on the GPU. KVCacheManagerV2 encodes a
 23Page as NVFP4 only while the Page resides in Host or Disk storage.
 24
 25```bash
 26python llm_kv_cache_compression.py \
 27    --compression-method nvfp4-cold-page \
 28    --model Qwen/Qwen3.5-4B
 29```
 30
 31A short request can remain entirely on the GPU. Use a workload with enough KV
 32pressure to trigger offload when validating the cold-page codec path.
 33
 34TriAttention
 35------------
 36TriAttention periodically evicts less important decode tokens. It requires an
 37offline calibration file produced for the selected model.
 38
 39```bash
 40python llm_kv_cache_compression.py \
 41    --compression-method triattention \
 42    --model Qwen/Qwen3-8B \
 43    --calibration-path /path/to/qwen3-8b-calibration.pt
 44```
 45"""
 46
 47import argparse
 48
 49from tensorrt_llm import LLM, SamplingParams
 50from tensorrt_llm.llmapi import (
 51    ColdPageQuantizationCompressionConfig,
 52    KvCacheConfig,
 53    TriAttentionKvCacheCompressionConfig,
 54)
 55
 56_DEFAULT_MODELS = {
 57    "nvfp4-cold-page": "Qwen/Qwen3.5-4B",
 58    "triattention": "Qwen/Qwen3-8B",
 59}
 60
 61
 62def _generate(
 63    model: str,
 64    kv_cache_config: KvCacheConfig,
 65    compression_config: ColdPageQuantizationCompressionConfig
 66    | TriAttentionKvCacheCompressionConfig,
 67) -> None:
 68    with LLM(
 69        model=model,
 70        backend="pytorch",
 71        trust_remote_code=True,
 72        max_seq_len=4096,
 73        max_batch_size=4,
 74        kv_cache_config=kv_cache_config,
 75        kv_cache_compression_config=compression_config,
 76    ) as llm:
 77        outputs = llm.generate(
 78            ["Explain why prefix caching helps agentic workloads."],
 79            SamplingParams(max_tokens=128, temperature=0.0),
 80        )
 81        print(outputs[0].outputs[0].text)
 82
 83
 84def run_nvfp4_cold_page(model: str) -> None:
 85    """Keep active KV unchanged and quantize only Host/Disk Pages."""
 86    _generate(
 87        model,
 88        KvCacheConfig(
 89            use_kv_cache_manager_v2=True,
 90            dtype="auto",
 91            host_cache_size=8 << 30,
 92        ),
 93        ColdPageQuantizationCompressionConfig(quant="nvfp4"),
 94    )
 95
 96
 97def run_triattention(model: str, calibration_path: str) -> None:
 98    """Periodically compact decode KV using offline calibration."""
 99    _generate(
100        model,
101        KvCacheConfig(
102            use_kv_cache_manager_v2=True,
103            enable_block_reuse=True,
104            dtype="auto",
105        ),
106        TriAttentionKvCacheCompressionConfig(
107            budget=64,
108            beta=32,
109            eviction_mode="union",
110            calibration_path=calibration_path,
111        ),
112    )
113
114
115def parse_arguments() -> argparse.Namespace:
116    parser = argparse.ArgumentParser()
117    parser.add_argument(
118        "--compression-method",
119        choices=tuple(_DEFAULT_MODELS),
120        default="nvfp4-cold-page",
121    )
122    parser.add_argument(
123        "--model",
124        default=None,
125        help="Model path or Hugging Face ID. Each method has a default model.",
126    )
127    parser.add_argument(
128        "--calibration-path",
129        default=None,
130        help="TriAttention calibration .pt produced for the selected model.",
131    )
132    args = parser.parse_args()
133    if args.compression_method == "triattention" and not args.calibration_path:
134        parser.error("--calibration-path is required for TriAttention")
135    return args
136
137
138def main() -> None:
139    args = parse_arguments()
140    model = args.model or _DEFAULT_MODELS[args.compression_method]
141    if args.compression_method == "nvfp4-cold-page":
142        run_nvfp4_cold_page(model)
143    else:
144        run_triattention(model, args.calibration_path)
145
146
147if __name__ == "__main__":
148    main()