KV Cache 复用 (又名：前缀缓存)#

如何使用#

通过将环境变量 NIM_ENABLE_KV_CACHE_REUSE 设置为 1 启用。有关更多信息，请参阅配置文档。

何时使用#

在超过 90% 的初始提示在多个请求中相同（仅在最后的 token 上有所不同）的场景中，实施键值缓存可以显著提高推理速度。 这种方法利用了提示的高度相似性，从而可以有效地重用计算资源，并最大限度地减少最终变体的处理时间。

例如，当用户询问有关大型文档的问题时，大型文档在请求中重复，但提示末尾的问题却不同。 当启用此功能后，首次令牌时间 (TTFT) 通常会加速约 2 倍。

示例

  • 大型表格输入，后跟关于该表格的问题

  • 相同的大型表格输入，后跟关于该表格的不同问题

  • 相同的大型表格输入，后跟关于该表格的不同问题

  • 等等…

KV Cache 复用将从第二个请求及后续请求开始加速 TTFT。

您可以使用以下脚本来演示加速效果

import time
import requests
import json

# Define your model endpoint URL
API_URL = "http://0.0.0.0:8000/v1/chat/completions"

# Function to send a request to the API and return the response time
def send_request(model, messages, max_tokens=15):
    data = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
        "top_p": 1,
        "frequency_penalty": 1.0
    }

    headers = {
        "accept": "application/json",
        "Content-Type": "application/json"
    }

    start_time = time.time()
    response = requests.post(API_URL, headers=headers, data=json.dumps(data))
    end_time = time.time()

    output = response.json()
    print(f"Output: {output['choices'][0]['message']['content']}")
    print(f"Generation time: {end_time - start_time:.4f} seconds")
    return end_time - start_time

# Test function demonstrating caching with a long prompt
def test_prefix_caching():
    model = "your_model_name_here"

    # Long document to simulate complex input
    LONG_PROMPT = """# Table of People\n""" + \
    "| ID  | Name          | Age | Occupation    | Country       |\n" + \
    "|-----|---------------|-----|---------------|---------------|\n" + \
    "| 1   | John Doe      | 29  | Engineer      | USA           |\n" + \
    "| 2   | Jane Smith    | 34  | Doctor        | Canada        |\n" * 50  # Replicating rows to make the table long

    # First query (no caching)
    messages_1 = [{"role": "user", "content": LONG_PROMPT + "Question: What is the age of John Doe?"}]
    print("\nFirst query (no caching):")
    send_request(model, messages_1)

    # Second query (prefix caching enabled)
    messages_2 = [{"role": "user", "content": LONG_PROMPT + "Question: What is the occupation of Jane Smith?"}]
    print("\nSecond query (with prefix caching):")
    send_request(model, messages_2)

if __name__ == "__main__":
    test_prefix_caching()
在本页中