KV cache 重用#

本文档描述了 KV cache 页面如何被以相同提示开始的请求共享和重用。这可以显著降低首次 token 延迟，即生成第一个输出 token 所需的时间。许多用例都可以从中受益，包括多轮请求和系统提示。

如何启用 KV cache 重用#

启用 KV cache 重用有两个步骤。

  1. 模型必须支持它

KV cache 重用要求模型使用分页上下文注意力（paged context attention）构建。这可以通过 trtllm-build 完成

trtllm-build --use_paged_context_fmha enable

  1. KV cache 重用在 KVCacheManager 中默认启用

如果您正在运行 gptManagerBenchmark 应用程序，可以使用命令行开关禁用 KV cache 重用

gptManagerBenchmark --enable_kv_cache_reuse enable=false

如果您正在运行 Triton 服务器，可以通过参数启用 KV cache 重用

parameters: {
  key: "enable_kv_cache_reuse"
  value: {
    string_value: "true"
  }
}

如果您正在使用 Executor API 编写自己的应用程序，可以在创建 KvCacheConfig 对象时包含 enableBlockReuse=true 来启用 KV cache 重用。请注意，这是默认设置，如果您希望禁用 KV cache 重用，请传入 enableBlockReuse=false。

GptSession 计划被弃用，并且不支持 KV cache 重用。

为 P-tuning 启用 KV cache 重用#

使用 P-tuning 时，不同的请求可能使用相同的伪输入 ID（即值大于词汇大小的提示 ID）。这可能导致错误的 KV cache 重用，因为 TRT-LLM 无法仅通过输入 ID 来区分这些请求。为了正确地为 P-tuning 启用 KV cache 重用，用户应为每个输入 ID 提供一个额外的 ID (uint64)。正常输入 ID（即文本 token ID）的额外 ID 应始终为 0，而伪输入 ID 应具有大于 0 的额外 ID。使用相同提示嵌入（prompt embedding）的请求应使用相同的额外 ID，而使用不同提示嵌入的请求应使用不同的额外 ID。

示例：假设词汇大小为 100，这意味着正常的文本 token ID 范围是 [0, 99]，提示 ID 从 100 开始。

# Request 1 uses prompt embedding table 1
input_ids = [100, 101, 102, 103, 1, 2, 3, 4]
extra_ids = [1,   1,   1,   1,   0, 0, 0, 0]

# Request 2 uses prompt embedding table 2
input_ids = [100, 101, 102, 103, 1, 2, 3, 4]
extra_ids = [2,   2,   2,   2,   0, 0, 0, 0]

# Request 3 uses prompt embedding table 1 and different text tokens
input_ids = [100, 101, 102, 103, 5, 6, 7, 8]
extra_ids = [1,   1,   1,   1,   0, 0, 0, 0]

性能预期#

当两个请求以相同的局部提示开始时，可以重用 KV cache 状态。这减少了首次 token 延迟，即生成第一个输出 token 所需的时间。共享提示越长，相对于总提示长度，节省的效果越大。当两个相同的请求连续运行时，节省的效果最大，在这种情况下，第一个输出 token 的延迟接近后续 token 的延迟。

可能阻止 KV cache 重用的情况#

在可能发生 KV cache 重用时，有一些陷阱可能会阻止它。KV cache 状态只有在计算该状态的请求终止后才变得可重用。如果您有一个共享的系统提示，第一个请求将计算系统提示的 KV cache 状态，第二个请求将重用它，但这仅在第二个请求在前一个请求完成后启动时才发生。如果使用大批量大小运行，很可能许多共享相同系统提示的请求会在第一个请求终止之前启动。在这种情况下，直到其中一个请求终止后，后续调度的请求才能重用 KV cache。

系统提示的 KV cache 状态将一直保持可重用，直到启动新请求或推进现有请求需要内存为止。发生这种情况时，可重用块会基于 LRU 策略被逐出。频繁使用的系统提示更有可能保持可重用，但不能保证，因为启动新请求的优先级高于可能的重用。例如，使用更大的批量大小或更大的输出序列长度会增加内存需求，从而降低 KV cache 块被重用的概率。

KV cache 状态存储在块中，每个块包含多个 token。只有完整的块才能被多个请求共享，因此块大小很重要。块大小是一个权衡：更大的块大小可能提高计算内核的效率，但会降低 KV cache 状态重用的可能性。块的默认大小为 128 个 token，这可以在使用 trtllm-build 命令构建模型时更改，例如

trtllm-build --tokens_per_block 32 ...

将创建一个模型，其中一个 KV cache 块可以容纳 32 个 token。请注意，tokens_per_block 必须是 2 的幂。

卸载到主机内存#

卸载到主机内存增加了 KV cache 重用的可能性。对于优先级更高的任务（如推进已运行的请求）所需的、可重用的块会被复制到主机内存的缓冲区中，而不是被逐出。这极大地扩展了可用于重用的内存量，使得块可以更长时间地保持可重用。另一方面，卸载块（以及在块被重用时重新加载）会产生一些开销，因为块必须在 CPU 和 GPU 内存之间来回复制。这种开销在 Grace-Hopper 机器上可以忽略不计，在配备 Hopper GPU 的 x86 机器上的许多用例中也足够小，可以产生净收益。由于 GPU 和主机内存之间的链接相对较慢，卸载在较旧的架构上不太可能带来收益。

如果您正在运行 gptManagerBenchmark，可以使用命令行开关启用卸载。例如，

gptManagerBenchmark --kv_host_cache_bytes 45000000000

将在主机内存中创建一个 45 GiB 的卸载缓冲区。请注意，此缓冲区是锁定内存（pinned memory），在 x86 机器上分配大量锁定内存可能需要相当长的时间（数十秒）。这是一次性开销。

如果您正在运行 Triton 服务器，可以使用 kv_cache_host_memory_bytes 参数启用卸载到主机内存。例如，将此参数添加到模型配置文件中将在主机内存中创建 45 GiB 的卸载缓冲区。

parameters: {
  key: "kv_cache_host_memory_bytes"
  value: {
    string_value: "45000000000"
  }
}

如果您正在使用 Executor API 编写自己的应用程序，可以在创建 KvCacheConfig 对象时包含 hostCacheSize=45000000000 来启用卸载到主机内存。这将在主机内存中创建 45 GiB 的卸载缓冲区。

GptSession 计划被弃用，并且不支持 KV cache 块卸载。