cannot broadcast [5, 5] to [1, 9, 5, 29]
主线位置：推理主线第四步「带 KV 缓存生成」，照着 transformers 的手写循环搬：每步只送一个 token、不写位置；一个缓存对象跑完一段对话接着用。
版本：PyTorch 2.15.0 / transformers 5.22.0（CPU）、candle 0.11.0（CPU，accelerate）
flowchart TB
f["forward"] --> rope["RoPE 按位置取行"]
f --> cat["k v 接在缓存后面"]
rope --> p0{"index_pos 对吗"}
p0 -->|"一直 0"| g1["静默乱位置"]
cat --> old{"缓存是旧对话的吗"}
old -->|"多 token"| e1["cannot broadcast"]
old -->|"单 token"| g2["静默串题"]
实验代码：PyTorch 侧 kv.py，candle 侧 crate pitb-exp 的 bin p21_cache（release）。模型是本地的 SmolLM2-135M-Instruct 权重，f32；PyTorch 固定 torch.set_num_threads(8)、注意力用 eager。prompt 是 The capital of France is（5 个 token），第二段是 The capital of Germany is。两边的缓存对象逐项对照见 KV 缓存；quantized_llama 上同一类错误见源码卷的坑 cannot broadcast [13, 45] to [1, 9, 13, 13]。
PyTorch 的写法
transformers 的手写贪心循环里没有位置，一个 DynamicCache 从头用到尾：
def greedy(prompt, cache, n=20):
out = model(prompt, past_key_values=cache, use_cache=True)
nxt = out.logits[0, -1].argmax().view(1, 1)
toks = [nxt.item()]
for _ in range(n - 1):
out = model(nxt, past_key_values=out.past_key_values, use_cache=True)
nxt = out.logits[0, -1].argmax().view(1, 1)
toks.append(nxt.item())
return toks
cache = DynamicCache(config=model.config)
first = greedy(ids, cache)
again = greedy(second, cache) # 不清缓存，接着用
fresh = greedy(second, DynamicCache(config=model.config))
cache type DynamicCache layers 30 seq_length 24
layer 0 keys (1, 3, 24, 64) torch.float32
first: ' Paris. Paris is the largest city in France and the capital of the French department of the Espace'
reuse: seq_length 48 ' Berlin. Berlin is the largest city in Germany and the capital of the German state of the EThe'
fresh: ' Berlin. Berlin is the largest city in Germany and the capital of the Federal Republic of Germany. Berlin'
位置是 LlamaModel.forward 用 past_key_values.get_seq_length() 推出来的，所以循环里不用管。复用缓存不报错：第二段接在第一段的 24 个位置后面，从第 15 个 token 起和新缓存的结果不同（ German state 对 Federal Republic）。这在 transformers 里也是错的，只是同样安静。
搬到 candle 的现象
candle 的 Llama::forward(x, index_pos, cache) 多一个参数。照 Python 的样子不管位置，decode 时一直传 0（p21_cache pos0），和每步传正确位置（ok）比：
pos0: " Paris, the capital of France, the United States, the United States, and the capital of France"
ok: " Paris. Paris is the largest city in France and the capital of the French department of the Espace"
不报错，第二个 token 起就变了。ok 和 transformers 逐个相同。
复用同一个 Cache 跑第二个 prompt，index_pos 从 0 开始（reuse）：
Error: cannot broadcast [5, 5] to [1, 9, 5, 29]
第二个 prompt 只有一个 token The 时（reuse1）：
reuse1: " capital of France. It is the largest city in Europe by population and the second most populous department of"
这次不报错，生成的是接着法国那一段往下说的话。
原因
CausalSelfAttention::forward 对 index_pos 只做两件事，都不和缓存的长度比对：
apply_rotary_emb用它从 cos / sin 表里取行：cache.cos.narrow(0, index_pos, seq_len)。一直传 0，每个新 token 都按位置 0 旋转，和缓存里键之间的相对位置全错。Cache::mask用它定 mask 的列数index_pos + seq_len。
而键和值总是接到缓存后面：
if cache.use_kv_cache {
if let Some((cache_k, cache_v)) = &cache.kvs[block_idx] {
k = Tensor::cat(&[cache_k, &k], 2)?.contiguous()?;
复用缓存、index_pos 为 0、5 个 token：mask 按 5 × 5 建，键却有 24 + 5 = 29 个，(1, 9, 5, 29) 的分数广播不上。单 token 时 seq_len == 1 不建 mask：
let att = if seq_len == 1 {
att
} else {
let mask = cache.mask(seq_len, index_pos)?.broadcast_as(att.shape())?;
masked_fill(&att, &mask, f32::NEG_INFINITY)?
};
没有形状要对，The 看得见缓存里关于法国的 24 个位置。llama.rs 的 Cache 没有清空的方法，index_pos == 0 时也不丢旧缓存（quantized_llama 会丢，见 KV 缓存）。所以 candle 这边，同一个错误在多 token 时报错、单 token 时静默；transformers 两种都静默。origin/main 上 llama.rs 和 0.11.0 相同。
改法
位置由调用方记着：prefill 传 0，之后每步加上刚送进去的 token 数。每段新对话新建一个 Cache：
let mut pos = start + ids.len();
for _ in 1..n {
let x = Tensor::new(&[next], dev)?.unsqueeze(0)?;
next = model.forward(&x, pos, cache)?.squeeze(0)?.argmax(0)?.to_scalar::<u32>()?;
out.push(next);
pos += 1;
}
// 新对话
let mut cache = Cache::new(true, DType::F32, &cfg, &dev)?;
fresh: " Berlin. Berlin is the largest city in Germany and the capital of the Federal Republic of Germany. Berlin"
fresh1: "\n\nThe key to a successful online course is to create a clear and engaging learning experience that resonates"
fresh 和 transformers 新缓存的那一行逐字相同。fresh1 是单个 The 在新缓存上的续写，和上面串题的 reuse1 完全不同。分段 prefill（seq_len > 1 且 index_pos > 0）是允许的，index_pos 写对就行（KV 缓存）。
怎么检查
- 把
Cache和位置包在一个自己的结构体里，forward只接 token，位置在结构体里加；Cache字段私有，没法从它本身读出长度，只能自己记。 - 对照测试：同一段 token，带缓存逐个送和不带缓存（
Cache::new(false, ..)）一次送，最后一个位置的 logits 应在 1e-4 以内；位置写错时这一条立刻不过。 - 新对话的入口处新建
Cache，不要靠报错提醒：单 token 的输入不会报错。
小结
- transformers 从缓存长度推位置；candle
llama.rs的index_pos由调用方传，不核对。一直传 0 不报错，生成从第二个 token 起变样。 - 复用
Cache跑新 prompt：多 token 报cannot broadcast [5, 5] to [1, 9, 5, 29]，单 token 静默接着上一段说；transformers 两种都静默接上。 - 改法：自己累加
pos，每段新对话Cache::new；改后和 transformers 逐个相同。