KV cache因果掩码:偏移怎么对齐
带KV cache的因果掩码,不能总按当前query张量从零开始画下三角。query可能只包含新来的一个或几个token,key却包含历史与当前token;两条轴的起点不同。先判断:这个键的逻辑位置,是否不晚于当前query的位置。
文章目录
下面用一个两行五列的反例,把错误掩码产生的数值偏差算出来,再给出可以独立运行的参考检查。Read in English。
先给两条轴标上真正的位置
设历史已经缓存P个token,本次输入L个新token,采用从零开始的连续逻辑编号。query行i的位置为P+i,拼接后的key列j的位置为j,键总长度S=P+L。对于本篇标准自注意力例子,允许访问的条件就是j≤P+i。
当没有历史缓存,P=0,条件退化为j≤i,熟悉的左上角下三角正好适用。带缓存以后公式没有变复杂,只是query行的起点变了。把“局部行号”当成“全局逻辑位置”,才会丢掉那段偏移。
Transformers缓存说明也要求迭代生成时的掩码覆盖过去与当前键。站内KV cache基础文章讨论缓存保存什么;这里专门处理这些键对当前query是否可见。
历史 P=3
K位置0、1、2
当前 L=2
Q位置3、4
拼接后的键
K位置0、1、2、3、4
逐行比较位置
key_position ≤ query_position
在键连续且query正好对应末尾L个位置的本例里,偏移也可写为S−L。但如果缓存经过窗口截取、含空槽或采用不同位置映射,不能仅凭矩阵尺寸恢复逻辑坐标。位置比较是通用的检查方法,几何对齐只是满足特定前提时的简写。
两行五列,正确答案是什么
取P=3、L=2。第一行query位于3,应看见key 0、1、2、3,不能看见未来的4;第二行query位于4,应看见全部五个键。可见键数依次是4和5。
若直接对2×5矩阵取左上角下三角,第一行只保留列0,第二行只保留列0和1。形状依然是2×5,矩阵乘法也可能正常完成,但当前token需要的历史上下文已被错误切掉。
✓
✓
✓
✓
×
✓
✓
✓
✓
✓
单token也可能被错误地限制到第一个键
当L=1、S=P+1且键只含过去和当前有效token时,唯一query处在逻辑位置P,这一行可以访问所有这些键。因此在这个限定场景,因果条件本身不再需要排除其中任何一列。
这不等于“单token解码永远可以忽略掩码”。静态缓存中未写入的槽位、padding、滑动窗口与其他模型约束仍可能要求屏蔽。应先明确key张量里实际放了什么,再决定是否需要额外条件。
让错误掩码给出可区分的结果
设五个key的注意力logits都是0,value分别为10、20、30、40、50。对允许位置做softmax,它们权重相同,所以正确两行输出为前四个值的平均25与全部五个值的平均30。错误下三角却给出10与15。
这是为排错设计的构造输入:value故意不同,才能让少看了历史这件事体现在输出上。如果所有value都相同,两种错误与正确掩码可能恰巧得到相同结果,导致一个薄弱测试误判通过。
# Pure-Python reference: True means allowed. No model or GPU is used.
import math
def causal_mask(query_positions, key_positions, key_valid=None):
valid = [True]*len(key_positions) if key_valid is None else key_valid
if len(valid) != len(key_positions):
raise ValueError("key validity length mismatch")
return [[bool(ok) and kp <= qp for kp,ok in zip(key_positions,valid)]
for qp in query_positions]
def attend(logits, values, allowed):
if not (len(logits)==len(values)==len(allowed)):
raise ValueError("row length mismatch")
kept = [(x,v) for x,v,ok in zip(logits,values,allowed) if ok]
if not kept:
raise ValueError("no valid key for this query")
peak = max(x for x,_ in kept)
weights = [math.exp(x-peak) for x,_ in kept]
return sum(w*v for w,(_,v) in zip(weights,kept))/sum(weights)
keys = list(range(5))
correct = causal_mask([3,4],keys)
wrong = [[j <= i for j in keys] for i in range(2)]
assert correct == [[True,True,True,True,False], [True]*5]
values, logits = [10,20,30,40,50], [0.0]*5
assert [attend(logits,values,m) for m in correct] == [25.,30.]
assert [attend(logits,values,m) for m in wrong] == [10.,15.]
assert causal_mask([4],keys) == [[True]*5]
# Padding is independent of causality, even when its position is old.
assert causal_mask([4],keys,[False,True,True,True,True])[0][0] is False
try:
attend([0.],[10.],[False])
except ValueError:
pass
else:
raise AssertionError("all-masked row was accepted")
print("offset, single-token, padding, and empty-row checks passed")
代码已本地执行,验证分块偏移、单token可见性、padding叠加与全屏蔽行。它只运行Python标量参考逻辑,没有调用PyTorch、运行模型或测量GPU性能;结果25/30和10/15是精确构造例子的输出。
API里的causal,不一定采用你想象的对齐
PyTorch 2.14 SDPA文档说明,非方阵的is_causal使用左上角因果偏置。因此不能仅凭参数名,就把它理解成自动知道当前query来自缓存末尾。对于本例的后缀query,应确认接口如何表达所需位置关系。
若采用显式布尔attn_mask,这个接口中True表示允许参与注意力;文档也提醒,它与MultiheadAttention的key_padding_mask语义相反。两个都是布尔矩阵,不代表可以原样传递。先在小矩阵里打印允许位置,比观察自然语言输出更可靠。
本篇不提供未执行的库调用片段。实际接入时需同时核对所装版本、掩码广播维度、所用后端,以及显式掩码和is_causal的组合限制。参考逻辑是验收标准之一,不能代替对具体内核行为的检查。
padding和因果条件要分别表达
因果条件回答时间先后,padding有效性回答这一列是否代表有效token。一个padding键即使逻辑位置很早,也不应因此被允许。参考函数先比较位置,再与key_valid相交,刻意保留两项条件的区别。
批处理中不同序列的历史长度可能不同,不能用整个batch的最大长度替代每个样本的真实位置。左padding、已经淘汰的窗口位置或分页存储还会让物理下标更不直观。把每行对应的query逻辑位置和每列对应的key逻辑位置列出来,才能检查访问关系。
如果query行本身是padding,还要按模型约定处理它的输出或后续损失。给这一行硬塞一个可见key来避免报错,不会自动让这个query变成有效数据;异常处理不应悄悄改变语义。
把被屏蔽logit置零,不等于屏蔽
softmax中的零logit仍有正的指数权重。假设两项logits原本都是0、value为10和1000,第二项不允许访问。正确输出是10;若只是把第二项logit乘以零,它仍为0,两项各占一半,输出反而是505。这个反例已用同一标量函数另行核对。
参考实现直接从归一化集合里排除不允许的项。真实张量实现常通过合适的加性偏置表达排除,但仍要遵循接口数值约定。不要把布尔可见性矩阵与任意逐元素乘法混为一谈,尤其不能只看结果形状判断掩码已经生效。
全屏蔽行需要单独验收
没有可访问键时,softmax不是普通的有效概率归一化。不同数值路径可能出现NaN、零或特定约定结果,不能把其中任意一种行为当作所有实现的保证。教学函数选择明确报错,让测试先暴露条件不完整。
生产代码怎样处理,应由模型与接口约定决定。若需要允许这种行,就给它定义清楚的行为并写专门测试;不要把负无穷替换成任意大负数后,误以为全屏蔽行自然拥有了正确概率。
用完整前向对照分块前向
把同一组输入分别一次性计算与分块带缓存计算,对照相应query位置的输出,是有辨识度的集成检查。保持模型权重、位置编号、精度和其他条件一致;若评估路径仍有dropout,先确认关闭规则,否则随机性会掩盖或伪装差异。
从最小用例开始:P=0的普通三角、P>0且L=1、P>0且L>1,再加入padding、静态空槽和模型确实使用的窗口限制。每次只增加一个条件,失败后才知道是哪项语义没有对齐。
今天的RoPE位置验证检查向量使用了哪个位置,本文检查哪些键允许参与。这两项需要一致,却各自可能独立出错。若这些反例帮你少猜一次缓存解码故障,欢迎在文末自愿支持。
来源与测试范围
PyTorch 2.14 SDPA参考;Transformers缓存机制。图、标量函数和数值为原创教学内容。本篇验证掩码的逻辑参考,不宣称验证任何具体GPU内核、模型质量或推理速度。
支持
如果文中的解释或示例帮你解决了问题、节省了时间,欢迎自愿打赏支持本站。
二维码可点击放大。更多支持方式见支持页面。


