TL;DR
- 采用了 PagedAttention,可以有效管理 attention 的 keys、values
- 吞吐量最多可以达到 huggingface 实现的24倍,并且不需要对模型结构进行任何的改变
1、PagedAttention
背景:
- LLM 的推理,最大的瓶颈在于显存。
- 自回归模型的 keys 和 values 通常被称为 KV cache,这些 tensors 会存在 GPU 的显存中,用于生成下一个 token。
- 这些 KV cache 都很大,并且大小是动态变化的,难以预测。已有的系统中,由于显存碎片和过度预留,浪费了60%-80%的显存。
实现:
- 受到操作系统中,虚拟内存和分页经典思想的启发
- PagedAttention 允许在不连续的内存空间中存储连续的 keys 和 values。 具体来说,PagedAttention 会将每个序列的 KV cache 划分为块,每个块包含固定数量 tokens 的 keys 和 values。 在注意力计算过程中,PagedAttention 内核有效地识别并获取这些块。
- 分块之后,这些 KV