从 KV 元素到账本与准入预算

读者位置与空间口径

本章承接第二章的逐层缓存接口。需要区分三种同时存在的对象：权重在请求之间共享，KV 缓存随各请求历史增长，临时激活用于当前前向计算。训练时的参数梯度和优化器状态不属于本课推理账本，但若实际进程仍保留它们，就不能从进程内存中假装扣除。空间测量必须说明是逻辑张量字节、库分配器持有字节、进程驻留内存还是设备显存。它们的差异来自副本、分配粒度、缓存池和运行时，而非单位换算。

令层数为 $N$，同长度批量为 $B$，缓存长度为 $L$，键值头数为 $h_{kv}$，每头宽度为 $d_h$，每元素字节数为 $b$。每层有 K、V 两份，均含 $BLh_{kv}d_h$ 个元素，乘上层数与字节数即得下式。该式只适用于所有层相同头结构、未压缩的稠密缓存。若各请求长短不同，可把 $BL$ 换成实际长度之和；若实现把全部请求补齐到最大长度，则仍要按分配后的长度计算。

完整算例一：公式必须对得上实际张量

项目模型有两层、批量二、缓存长二十四、四个键值头、每头宽八，使用双精度每元素八字节。单份单层张量有 $2\times4\times24\times8=1536$ 个元素，占一万二千二百八十八字节。K、V 两份与两层共四份，总量四万九千一百五十二字节，恰为四十八 KiB。注意 KiB 按一千零二十四字节计算，kB 按一千字节计算；如果报四十九点一五 kB，也应写清十进制单位。项目通过 numel 乘 element_size 独立求和后与公式相等。

from lab import make_model,tokens,chunks,kv_bytes
model=make_model();_,past=chunks(model,tokens(24,2),[24])
actual=sum(t.numel()*t.element_size() for pair in past for t in pair)
expected=kv_bytes(2,2,24,4,8,8)
assert actual==expected==49152
print(actual,actual/1024)

零长度是一个有用的极限：逻辑缓存应为零，而模型权重仍在。批量从一翻到二时缓存翻倍，共享权重不翻倍。每元素从八字节改到两字节时同形状逻辑缓存缩为四分之一，但模型必须真以该 dtype 保存缓存才成立；只在报告里修改 b 不会改变已分配张量。还应检查每层 K、V 的实际 dtype，因为混合精度实现可能让归一化、权重与缓存采用不同精度。项目据实际双精度计账，没有冒充半精度推理。

GQA 不只改变缓存，还改变 K、V 投影

普通多头注意力有 $h_q=h_{kv}$，且 $d=h_qd_h$。分组查询注意力保留 $h_q$ 个查询头，但让多个查询头共享一组键值头，要求本课简单分组中 $h_q$ 能被 $h_{kv}$ 整除。查询和输出投影仍各含 $d^2$ 个权重；键和值投影各从宽度 d 输出到 $h_{kv}d_h$，各含 $dh_{kv}d_h$ 个权重。忽略偏置的注意力投影总数如下。减少键值头并未减少查询数，因此每个查询对可见位置的打分工作不会简单按键值头比例消失。

例如 $d=128,h_q=8,d_h=16$。普通多头的投影为六万五千五百三十六个；两键值头的分组查询为四万零九百六十个，比例是零点六二五，而缓存比例是四分之一。若误写“所有注意力参数也缩为四分之一”，会漏掉未变的 Q、O。前馈、嵌入和归一化参数还应另加，所以整模型参数缩减比例更小。项目实际模型仍是普通多头，本章的分组计算器只做账本，不把改头数当成已实现的输出等价优化。

把训练好的多个头简单平均到较少键值头，通常会改变函数，需要训练或校准与质量验证。原始 GQA 研究讨论过从已有多头检查点转换再继续训练的方案；本课不复现其训练结果，也不引用其硬件速度作为本实验成绩。对本项目而言，GQA 是一个明确的架构备选，属于需要另立检查点和质量证据的改变。缓存复用的精确证明不能证明改架构后精确等价，这两类改变应在最终报告分栏。

从逻辑长度到实际分配

连续预留若按最大长度分配，短请求会占住未来可能用不到的槽位。分页思想把每请求的逻辑词元位置映射到固定大小块，增长时按块取得空间。设一块容纳 $S$ 个词元，请求 i 当前长为 $L_i$，需块数为 $\lceil L_i/S\rceil$，总槽位为各块数之和乘 S。最后一块未用槽位是内部碎片；块表和引用计数还会额外占空间。逻辑不连续的物理块需要注意力内核支持映射访问，单纯计算 ceil 并不会让当前模型变成分页推理器。

完整算例二：块越小不代表系统越快

请求长度为十七、三十一、二，块大小十六，则块数分别二、二、一，总共五块八十槽，实际使用五十槽，浪费三十槽。若每槽 KV 成本一千字节，仅缓存载荷分配八万字节，其中三万字节未使用。块大小改为八，块数三、四、一，总六十四槽，浪费十四槽；但块数由五变八，块表与索引管理增加。此处只证明内部碎片减少，没有测到寻址开销、物理连续性或内核带宽，所以不能推断块八必快。

安全准入还要为未来增长留预算。设当前可用预算为 F，预留不可动用余量 R，每个将被接纳请求按约定最大长度估计需 K 字节缓存，则同质请求的简单容量上界为 $\max(0,\lfloor(F-R)/K\rfloor)$。这是假设权重等固定占用已扣除、请求独立占用且临时激活被余量覆盖的保守算术；如果当前 F 是设备总容量而没有扣权重，就会高估。若只按当前一词元占用准入，又没有增长限额，稍后会同时溢出。

取可用一千零二十四字节、预留六十四、每请求一百九十二字节，余量后剩九百六十，最多五个请求。若余量增到一百二十八，结果变成四个。这个不连续跳变来自向下取整，不能用“平均每请求占一百九十二”把小数容量解释为可接纳半个请求。异质长度应逐请求扣预算，且生成过程中应在增长前检查下一块是否能分配；失败时可排队、拒绝或按已声明策略抢占，不能继续写越界。

import math
lengths=[17,31,2]
def ledger(block):
    allocated=sum(math.ceil(n/block) for n in lengths)*block
    return allocated,allocated-sum(lengths)
assert ledger(16)==(80,30) and ledger(8)==(64,14)
assert (1024-64)//192==5 and (1024-128)//192==4
d,hq,hkv,dh=128,8,2,16
assert 2*d*d+2*d*hkv*dh==40960
print(ledger(16),ledger(8))

把峰值和生命周期放回账本

前向过程中会短暂同时持有旧缓存、拼接后的新缓存、打分矩阵和 logits。逻辑缓存求和只量最终返回张量，不能证明峰值等于该值。尤其当前实现显式构造注意力矩阵，其预填充临时量按长度平方增长；FlashAttention 采用重排计算减少这类中间物化和数据搬运，但不把模型的合法注意力关系改成近似稀疏关系。当前项目没有该内核，因此只把它列为可进一步研究的实现机制，不提供任何虚构的速度或峰值记录。

第一类检查是维度与单调性：长度、层数、批量、每元素字节翻倍时对应逻辑成本翻倍，键值头减少时缓存按比例变小。第二类检查是独立观察对象：实际张量求和对公式，分配槽位对真实长度，预算对剩余整数容量。这些检查分别约束不同层次，无法相互替代。只有真实服务运行中跟踪高水位、临时峰值和回收行为，才有资格承诺不会因特定请求混合而耗尽内存；本实验尚未实现该服务。

总参数也可以用实物清点而不是猜测“几百万模型”。逐项累加 model.parameters 的 numel 得到元素数，再按各自 element_size 求字节，能够发现绑定权重、重复引用与混合 dtype 带来的口径问题。本模型头部不与嵌入绑定，因而词表嵌入与输出投影分别占空间；若未来改为共享同一存储，应按唯一底层存储核对，而不是将参数名字出现次数直接相加。保存文件还可能有容器元数据和对齐，因此磁盘检查点大小与参数裸字节也不必相等。

练习与项目增量

项目增量是 kv_bytes、pages、capacity 三个可独立检验的账本函数，以及实际缓存字节证据。第四章将用有限槽位讨论调度，但那里的槽位是服务并行位置，不等于本章字节页块。两个模型相连时必须有从每个活跃请求到缓存预算的映射；当前实验把它们分开验证，避免将离散调度吞吐当作真实显存管理已经完成。