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 块卸载。