题目重点：缓存已经投影的K/V，每次只计算新token，但新chunk仍需要内部因果mask。
题意与假设
这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。
- 沿用单层causal decoder，无dropout，缓存为B×H×pastT×Dh的K/V；同batch、同参数、同前缀，训练没有cache淘汰或跨模型复用。
- 新输入可一个token也可一个chunk，绝对位置从pastT开始。缓存不保存token logits或完整历史hidden，且示例不支持padding变长cache。
解题路线
- 新tokens只算新Q/K/V，K/V与cache沿序列维concat；query位置offset+i仅能看key_index≤offset+i。单token无future，chunk多token仍有未来位置，不能把cached模式统统设无mask。
- 位置embedding使用p.position[offset:offset+t]，每次从0开始会让incremental输出与full forward不同。attention key长度past+t，scale仍√Dh，与past长度无关。
- cache需要绑定request、model revision、adapter和token prefix，不能按最后token简单复用。批处理不同有效长度需要position IDs和有效key mask，物理padded cache.shape不是每个用户真实position。
- 比对三个路径：完整T、逐token、任意切点prefix+suffix chunk，logits应在误差内一致；同时改变未来token验证因果性。拼接实现复制全部历史K/V，证明正确不等于高效生产paged cache。
手动走一遍
past4、新chunk3，query0可看key0..4、query1看0..5、query2看0..6；新chunk第一query不能看它后面的两token。同prompt、不同adapter产生不同K/V，因此cache命中必须包含adapter标识。请求取消后cache释放，但并发decode引用尚未结束时要引用计数/lease，不可提前复用内存页。
Python 参考实现
下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。
代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。
import numpy as np
def layer_norm(x):
return (x-x.mean(axis=-1,keepdims=True))/np.sqrt(x.var(axis=-1,keepdims=True)+1e-5)
def gelu(x):return .5*x*(1+np.tanh(np.sqrt(2/np.pi)*(x+.044715*x**3)))
def initialize(vocab,width,heads,max_length,seed=0):
if min(vocab,width,heads,max_length)<1 or width%heads:raise ValueError('dimensions/heads')
rng=np.random.default_rng(seed);scale=width**-.5
return dict(token=rng.normal(size=(vocab,width))*scale,position=rng.normal(size=(max_length,width))*scale,
q=rng.normal(size=(width,width))*scale,k=rng.normal(size=(width,width))*scale,
v=rng.normal(size=(width,width))*scale,o=rng.normal(size=(width,width))*scale,
ff1=rng.normal(size=(width,4*width))*scale,ff2=rng.normal(size=(4*width,width))*scale,heads=heads)
def forward(tokens,p,cache=None):
ids=np.asarray(tokens)
if ids.ndim!=2 or ids.shape[1]<1 or not np.issubdtype(ids.dtype,np.integer):raise ValueError('B,T integer tokens')
if np.any(ids<0) or np.any(ids>=len(p['token'])):raise ValueError('token ID')
b,t=ids.shape;d=p['token'].shape[1];h=p['heads'];dh=d//h
offset=0 if cache is None else cache['k'].shape[2]
if offset+t>len(p['position']):raise ValueError('position length')
x=p['token'][ids]+p['position'][offset:offset+t];normalized=layer_norm(x)
split=lambda a:a.reshape(b,t,h,dh).transpose(0,2,1,3)
q=split(normalized@p['q']);k=split(normalized@p['k']);v=split(normalized@p['v'])
if cache is not None:
if cache['k'].shape!=(b,h,offset,dh) or cache['v'].shape!=cache['k'].shape:raise ValueError('cache shape')
k=np.concatenate([cache['k'],k],axis=2);v=np.concatenate([cache['v'],v],axis=2)
scores=q@k.swapaxes(-1,-2)/np.sqrt(dh)
permitted=np.arange(k.shape[2])[None,:]<=offset+np.arange(t)[:,None]
scores=np.where(permitted[None,None,:,:],scores,-np.inf)
shifted=scores-scores.max(axis=-1,keepdims=True);weights=np.exp(shifted);weights/=weights.sum(axis=-1,keepdims=True)
attended=(weights@v).transpose(0,2,1,3).reshape(b,t,d)
hidden=x+attended@p['o'];hidden=hidden+gelu(layer_norm(hidden)@p['ff1'])@p['ff2']
logits=layer_norm(hidden)@p['token'].T
return logits,hidden,dict(k=k,v=v)
def classify(tokens,p,pooling_mask,head,bias):
# Encoder-like use is a different contract; this classifier retains causal hidden states.
_,hidden,_=forward(tokens,p);mask=np.asarray(pooling_mask,dtype=bool)
if mask.shape!=hidden.shape[:2] or np.any(mask.sum(axis=1)==0):raise ValueError('nonempty pooling mask')
# Mask must be right padding; masked suffix cannot influence earlier causal states.
if np.any(mask[:,1:] & ~mask[:,:-1]):raise ValueError('right padding')
pooled=(hidden*mask[:,:,None]).sum(axis=1)/mask.sum(axis=1,keepdims=True)
return pooled@head+bias
复杂度
单tokenattention O(BpastT·D)，新投影/FFN O(BD²)，cacheO(BpastT·D)。示例concat每步O(BpastT·D)复制，累计生成T可O(BT²D)copy；预分配/paged结构减少copy但不消除attention读历史成本。
易错点与追问
容易失分的边界
- cache offset错位是隐蔽错误，short T测试可能不明显。
- 缓存Q并重算旧query没有带来decode节省。
- 多tokenchunk不屏蔽未来会泄露后续输入。
继续深挖
- GQA/MQA怎样改变K/V heads及显存估算？
- 滑动窗口或RoPE cache必须保存哪些绝对位置与旋转状态？
