做一个 KV 缓存显存估算器

目标

将KV缓存内存公式实现为代码，通过计算确认GQA和量化的效果，计算出给定GPU中可能的最大同步请求数。

为什么重要

在部署前，总会有人问“A100 80GB一张卡能装入这个模型吗？”如果只计算权重，答案就会是错误的。7B模型的权重是fp16，14GB，但如果放在128K上下文中，配置为8，KV缓存是512GiB，也就是权重的36倍。如果不进行这个计算，在部署当天就会遇到OOM，OOM是当负载增加时出现的，所以是在最糟糕的时候发生的。而且这个计算器一旦创建，就会一直写入——在更改模型时，在增加上下文长度时，在审查量化时，max_num_seqs在确定时每次都需要。第6步中获得的最大同时请求数量是容量计划的起点。

阶段

/root/kv/kv.py在kv_bytes(layers, hidden, seq, batch, dtype_bytes, kv_groups=1)制作。基本公式是2 * layers * hidden * seq * batch * dtype_bytes是。
layers=32, hidden=4096, seq=4096, batch=1, dtype_bytes=2用计算出来的值/root/kv/base.txt在bytes=<정수> gib=<소수 둘째자리>写在罗上。gib必须是2.00。
seq=131072换成罗的值/root/kv/long.txt用同样的格式写。gib应该是64.00。
kv_groups=8应用/root/kv/gqa.txt写在上面。gib必须持续8.00（以128K为基准）。
dtype_bytes=1应用/root/kv/int8.txt写在上面。gib必须达到4.00（以128K、GQA-8为标准）。
/root/kv/fit.txt在gpu_gib=80 weights_gib=14 overhead_gib=6 kv_budget_gib=60 per_req_gib=<소수> max_concurrent=<정수>写。每个请求的基准是32K上下文、GQA-8、fp16。
/root/kv/kv_table.csv在context,batch,gib写入头和6行。上下文是4096、32768、131072，配置是1和8。值是没有GQA的fp16标准。

参考

公式的前2个是高度和值两对的意思。
GQA将键值头以组为单位共享，因此按组数进行划分。
加权量化（GPTQ、AWQ）和KV缓存量化是独立设置。
常见的错误1：只计算权重而忽略KV缓存——在长上下文中，KV占主导地位。
常见的错误2：混淆GiB和GB——这个练习以2^30为基准。

正式实现

/root/kv/kv.py在kv_bytes(layers, hidden, seq, batch, dtype_bytes, kv_groups=1)制作。基本公式是2 * layers * hidden * seq * batch * dtype_bytes是。

请理解前面的2是什么，然后输入。参数是层数、隐藏维度、序列长度、布局、dtype字节。

以7B 4K基准值进行验证

layers=32, hidden=4096, seq=4096, batch=1, dtype_bytes=2用计算出来的值/root/kv/base.txt在bytes=<정수> gib=<소수 둘째자리>写在罗上。gib必须是2.00。

有被广泛引用的主基准值。如果在这里正确，那么公式就正确了。

扩展为长语境

seq=131072换成罗的值/root/kv/long.txt用同样的格式写。gib应该是64.00。

只要将序列长度增加32倍就可以了。看看结果为什么会让长对话变得困难。

减少应用GQA

kv_groups=8应用/root/kv/gqa.txt写在上面。gib必须持续8.00（以128K为基准）。

键和值头共享到组中。想想有多少个组，分割点在公式的哪儿。

应用KV缓存量化

dtype_bytes=1应用/root/kv/int8.txt写在上面。gib必须达到4.00（以128K、GQA-8为标准）。

减少dtype字节。与权重量化是独立设置这一点很重要。

获取GPU中最大同时请求数

/root/kv/fit.txt在gpu_gib=80 weights_gib=14 overhead_gib=6 kv_budget_gib=60 per_req_gib=<소수> max_concurrent=<정수>写。每个请求的基准是32K上下文、GQA-8、fp16。

从总内存中减去加权和开销，并按请求大小进行划分。这个值是同步性上限。

生成剧本表

/root/kv/kv_table.csv在context,batch,gib写入头和6行。上下文是4096、32768、131072，配置是1和8。值是没有GQA的fp16标准。

通过改变上下文和布局来制作表格。以后在容量计划时，请制作成可以原样使用的形式。