TensorRT LLM--Paged KV Cache

技术出处：vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention | vLLM Blog

kv cache出现的动机

decoder推理中，对于每个输入的 prompt，在计算第一个 token 输出的时候，每个 token 的 attention 肯定是都要从头计算, 但是在后续 token 的生成中，需要concat前面每一个 token 的 K 和 V，由于模型参数矩阵是不变的，此时只有刚生成的那个 token 的 K 和 V 需要从头计算，所以可以把之前token的K和V缓存起来避免重复计算，这个就叫kv cache

在生成阶段，一个常见的优化是为MHA内核提供一个缓存，该缓存包含已经计算的过去K和V元素的值，该缓存被称为KV缓存，TensorRT LLM使用该技术来加速其生成阶段。在TensorRT LLM中，每个Transformer层有一个KV缓存，这意味着一个模型中有与层一样多的KV缓存。TensorRTLLM的当前版本支持两种不同类型的KV缓存：连续KV缓存和分页KV缓存。

连续KV缓存

连续的KV缓存是一个整体张量。其形状为：

[max_batch_size * max_beam_width, 2, num_heads, max_seqlen, hidden_dim_per_head] 当序列比最大序列长度短时，该实现使用的内存要多得多。

分页KV缓存

分页KV缓存将KV缓存分解为块，这些块在处理过程中由高速缓存管理器分配给不同的请求。该缓存管理器跟踪序列，从池中分配新块，并在需要时回收这些块。请参阅TensorRT-LLM/tensorrt_llm/runtime/kv_cache_manager.py at release/0.5.0 · NVIDIA/TensorRT-LLM (github.com)的简化实现。Batch Manager中包含了更高效的C++实现。

分页KV缓存（paged attention）出现动机

虽然kv cache很重要，但是kv cache所占的空间也确实是大且有浪费的，所以出现了paged attention来解决浪费问题。kv cache大小取决于seqlen，然而这个东西对于每个batch里面的seq来说是变化的，毕竟不同的人输入不同长度的问题，模型有不同长度的答案回答，kv cache统一按照max seq len来申请，造成现有decoder推理系统浪费了很多显存。

分页KV缓存（paged attention）设计原理

Paged Attention的核心是一张表，类似于OS的page table，这里叫block table，记录每个seq的kv分布在哪个physical block上，通过把每个seq的kv cache划分为固定大小的physical block，每个block包含了每个句子某几个tokens的一部分kv，允许连续的kv可以不连续分布。在attention compute的时候，page dattention CUDA kernel就通过block table拿到对应的physical block序号，然后CUDA线程ID计算每个seq每个token的offset从而fetch相应的block，拿到kv，继续做attention的计算

分页KV缓存（paged attention）的工作流程

参考文章：(28 封私信 / 80 条消息) heyguy - 知乎 (zhihu.com)