KV缓存重用

本文档描述了如何通过以相同提示开头的请求共享和重用kv缓存页面。这可以大大降低首令牌延迟，即生成第一个输出令牌所需的时间。许多用例可以从中受益，包括多轮请求和系统提示。

如何启用kv缓存重用

启用kv缓存重用有两个步骤。

模型必须支持它

KV缓存重用要求模型为分页上下文注意力构建。这是通过trtllm-build完成的：

trtllm-build --use_paged_context_fmha enable

必须在KVCacheManager中启用KV缓存重用

如果您正在运行gptManagerBenchmark应用程序，您可以通过命令行开关启用kv缓存重用：

gptManagerBenchmark --enable_kv_cache_reuse enable

如果您正在运行Triton服务器，您可以通过一个参数启用kv缓存重用：

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

如果您正在使用Executor API编写自己的应用程序，您可以在创建KvCacheConfig对象时通过包含enableBlockReuse=true来启用kv缓存重用。

GptManager API 已被弃用，但如果您有一个使用 GptManager API 的旧应用程序，您可以通过一个可选参数启用 kv 缓存重用：

TrtGptModelOptionalParams 类封装了以下字段：
- kvCacheConfig
  - enableBlockReuse (默认: false) 允许跨请求重用先前计算的KV缓存块。这有望优化内存使用和计算。

GptSession 计划被弃用，并且不支持 kv 缓存重用。

启用kv缓存重用以进行p-tuning

在使用p-tuning时，不同的请求可能会使用相同的假输入ID（即提示ID的值大于词汇表大小）。这可能导致不正确的kv缓存重用，因为TRT-LLM无法仅通过输入ID来区分这些请求。为了正确启用p-tuning的kv缓存重用，用户应为每个输入ID提供一个额外的ID（uint64）。普通输入ID（即文本标记ID）的额外ID应始终为0，而假输入ID的额外ID应大于0。使用相同提示嵌入的请求应使用相同的额外ID，而使用不同提示嵌入的请求应使用不同的额外ID。

示例：假设词汇表大小为100，这意味着普通文本的标记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缓存状态可以被重用。这减少了第一个输出令牌生成所需的时间，即第一个输出令牌的延迟。当共享提示相对于整个提示长度较长时，可以实现更大的节省。当两个相同的请求连续运行时，可以实现最大的节省，在这种情况下，第一个输出令牌的延迟接近后续令牌的延迟。

可能阻止kv缓存重用的情况

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

系统提示的Kv缓存状态将保持可重用，直到需要内存来启动新请求或传播现有请求。当这种情况发生时，可重用块将基于LRU被驱逐。经常使用的系统提示有更好的机会保持可重用，但不能保证，因为启动新请求优先于可能的重用。例如，使用更大的批量大小或更长的输出序列长度运行将减少kv缓存块被重用的概率，因为它增加了内存需求。

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

trtllm-build --tokens_per_block 32 ...

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

卸载到主机内存

卸载到主机内存增加了kv缓存重用的可能性。对于更高优先级任务（如传播已经运行的请求）所需的可重用块，会被复制到主机内存中的缓冲区，而不是被驱逐。这大大扩展了可用于重用的内存量，使块能够保持更长时间的可重用性。另一方面，块的卸载（以及在块被重用时的重新加载）会有一些成本，因为块必须从CPU复制到GPU内存，反之亦然。在Grace-Hopper机器上，这种成本可以忽略不计，并且在配备Hopper GPU的x86机器上，对于许多用例来说，成本足够小，可以产生净收益。在较旧的架构上，由于GPU和主机内存之间的（相对）慢速连接，卸载不太可能产生好处。

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

gptManagerBenchmark --kv_host_cache_bytes 45000000000

将在主机内存中创建一个45 GiB的卸载缓冲区。请注意，此缓冲区是固定内存，在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的卸载缓冲区。

GptManager API 已被弃用，但如果您有一个现有的应用程序正在使用 GptManager API，您可以通过一个可选参数启用卸载：

TrtGptModelOptionalParams 类封装了以下字段：
- kvCacheConfig
  - hostCacheSize (默认值: 0) 用于在从GPU内存中逐出时卸载kv缓存页面的主机缓冲区的大小（以字节为单位）。

GptSession 计划被废弃，并且不支持 kv 缓存块卸载。