GPU LLM 内存需求的两个主要贡献者是模型权重和 KV 缓存。
- 模型权重：模型参数占用内存。例如，具有 70 亿个参数（例如 Llama-2-7b），以 16 位精度（FP16 或 BF16）加载时，显存大小约为 70 亿 * 2 字节（FP16）= 14 GB。
- KV 缓存：内存中的自注意力张量被用作缓存，以避免重复计算。
在批处理时，批量中每个请求的 KV 缓存仍必须单独分配，并且可能会占用大量内存。以下公式描述了适用于当今大多数常见 LLM 架构的 KV 缓存的大小。
每个令牌的 KV 缓存大小（字节）= 2 * (num_layers) * (num_heads * dim_head) * precision_in_bytes
第一个系数为 2 解释了 K 和 V 矩阵。通常，(num_heads*dim_head) 的值与 Transformer 的 hidden_size (或模型维度 d_model)相同。这些模型属性通常可在模型卡或相关配置文件中找到。
输入序列中的每批输入中的每个令牌都需要此内存大小。假设半精度，KV 缓存的总大小由以下公式给出。
以字节为单位的 KV 缓存总大小 = (batch_size) + (sequence_length) * 2 * (num_layers) * (hidden_size) * sizeof(FP16)
例如，对于 16 位精度的 Lama 2 7B 模型，批量大小为 1，则 KV 缓存的大小为 1*4096*2*32*4096*2 字节，即小于 2 GB.
高效管理此 KV 缓存是一项艰巨的任务。随着批量大小和序列长度的线性增长，显存需求可以快速扩展。因此，它限制了可以提供的吞吐量，并对长上下文输入构成了挑战。
