cannot broadcast [13, 45] to [1, 9, 13, 13]

主线位置：推理主线的 forward(input, pos)。index_pos 决定 rotary 表从哪一行取、mask 按多长建、要不要丢掉旧缓存；缓存的真实长度要到 forward_attn 里 cat 完才知道。机制和更多传错的情形见量化 llama 的一次前向，缓存本身见 KV 缓存与注意力。代码在 candle-transformers 的 models/quantized_llama.rs。

flowchart TB
  t1["第 1 轮 pos 到 32"] --> clear["clear_kv_cache"]
  clear --> keep{"pos 归零了吗"}
  keep -->|否 多 token| err["mask 13x45 缓存 13"]
  err --> boom["cannot broadcast"]
  keep -->|否 单 token| same["相对位置不变 输出相同"]
  keep -->|是| ok["正常"]
  step["pos 每次只加 1"] --> garbage["第 1 轮就乱码"]

实验 crate exp-candle/pb-pit，bin chat，SmolLM2-135M Q8_0 GGUF，贪心生成，CPU（Metal 上报错原文相同），release，M2 Max、macOS 27、rustc 1.99.0。

现象

一个最简单的多轮对话封装：模型加一个位置计数器，每轮先送 prompt，再一个一个送生成的 token：

struct Chat {
  model: ModelWeights,
  pos: usize,
  // ...
}

impl Chat {
  fn turn(&mut self, ids: &[u32], n: usize) -> anyhow::Result<String> {
    let mut out = vec![];
    let mut input = ids.to_vec();
    for _ in 0..n {
      let x = Tensor::new(input.as_slice(), &self.dev)?.unsqueeze(0)?;
      let logits = self.model.forward(&x, self.pos)?;
      self.pos += input.len();
      let next = logits.squeeze(0)?.argmax(0)?.to_scalar::<u32>()?;
      out.push(next);
      input = vec![next];
    }
    self.tok.decode(&out, true).map_err(E::msg)
  }
}

第一轮问「What is Rust?」，然后开始一段新对话：清缓存，问「What is Python?」，但忘了把 pos 归零：

chat.model.clear_kv_cache();
let p2 = prompt(&tok, "What is Python?")?;
println!("turn 2: {:?}", chat.turn(&p2, 20)?);

turn 1 (13 prompt tokens): "Rust is a modern, high-performance, and object-oriented programming language developed by Brian K"
pos after turn 1 = 32
turn 2 starts at pos 32 with 13 tokens
Error: cannot broadcast [13, 45] to [1, 9, 13, 13]
exit=1

报错说的是广播，不提 index_pos，也不提缓存。

另一种常见写法是计数器按「调用次数」而不是「token 数」前进（self.pos += 1，step 模式）。这回没有错误：

turn 1 (13 prompt tokens): "R\nRust is a programming languagesystem\n\nuser\nI am trying"
pos after turn 1 = 20
turn 2 starts at pos 20 with 13 tokens
Error: cannot broadcast [13, 33] to [1, 9, 13, 13]
exit=1

第一轮静默地生成了乱码（正确的输出见上一段的第一行），要到第二轮的多 token 输入才报错；如果程序只做单轮生成，就永远只有乱码。

出处

forward 只在多 token 输入时建 mask，长度由 index_pos 推出：mask 按 (seq_len, index_pos + seq_len) 建，这里是 13 × (32 + 13) = 13 × 45。缓存是另一回事，在 forward_attn 里：

        let (k, v) = match &self.kv_cache {
            None => (k, v),
            Some((k_cache, v_cache)) => {
                if index_pos == 0 {
                    (k, v)
                } else {
                    let k = Tensor::cat(&[k_cache, &k], 2)?;
                    let v = Tensor::cat(&[v_cache, &v], 2)?;
                    (k, v)
                }
            }
        };

clear_kv_cache 把缓存设成 None，于是这一轮的 k 只有 13 个位置，注意力分数是 (1, 9, 13, 13)；13 × 45 的 mask 广播不上去，报错的是 mask.broadcast_as(att.shape())。

step 模式下第一轮 prefill 之后 pos 是 1，而缓存里有 13 个位置：每个 decode 步的 rotary 用的是 1、2、3……，缓存里的键却在 0 到 12、之后接着增长。单 token 输入不建 mask，没有任何地方会发现两者对不上，注意力按错误的相对位置算，输出就坏了。

两个 forward 的参数，index_pos 是唯一需要调用方自己维护、而库不核对的状态。不变式只有一条：index_pos 等于缓存里已有的位置数。

📝NOTE

没有蠢问题

问：清了缓存却不归零，如果第二段对话的第一次输入只有一个 token（没有 mask），会生成乱码吗？

答：不会，至少在这里不会。实验里第二段只送一个「Python」token，pos 从 32 开始和从 0 开始，CPU 上生成的 20 个 token 逐字相同（Metal 上两者也相同）：

turn 2 starts at pos 32 with 1 tokens
turn 2: "\n\n```python\ndef find_duplicates(nums):\n    seen = set()\n    duplicates"
turn 2 starts at pos 0 with 1 tokens
turn 2: "\n\n```python\ndef find_duplicates(nums):\n    seen = set()\n    duplicates"

RoPE 让注意力分数只依赖查询和键的位置差。缓存清空之后，新对话里所有位置都整体平移 32，两两之差不变，所以结果一样（只差浮点舍入）。乱码出现在位置差被弄错的时候，也就是同一段对话里计数器和缓存脱节（上面的 step）。整体平移的代价是提前用掉 rotary 表：它只有 4096 行，index_pos + seq_len 超过就报 narrow 错误。

修法

计数器和缓存要一起变。把「开始新对话」写成一个方法，让它同时清缓存和归零；计数器按送进去的 token 数前进：

impl Chat {
  fn reset(&mut self) {
    self.model.clear_kv_cache();
    self.pos = 0;
  }
}

（实验里的 fix 模式就是在 clear_kv_cache() 之后加 chat.pos = 0。）

turn 1 (13 prompt tokens): "Rust is a modern, high-performance, and object-oriented programming language developed by Brian K"
pos after turn 1 = 32
turn 2 starts at pos 0 with 13 tokens
turn 2: "Python is a high-level, interpreted programming language developed by Guido van Rossum in 1"
exit=0

quantized_llama 在 index_pos == 0 时自己丢掉旧缓存，所以对它来说「归零」比「清缓存」更要紧：只归零不清缓存，结果也和新模型一样（量化 llama 的一次前向的动脑筋）；clear_kv_cache 的作用是提前释放 30 层的缓存张量。

怎么确认没中招

pos 只在一个地方前进，前进的量是 input.len()；第一轮生成完打印一次，应当等于 prompt 长度加生成的 token 数减 1（最后一个生成的 token 还没送进去）。上面是 13 + 20 − 1 = 32。
报 cannot broadcast [s, k] to [1, h, s, c] 时，k − s 就是你传的 index_pos，c − s 是缓存里真实的位置数，两者应当相等。
第一轮的贪心输出和一个新加载的模型、pos 从 0 正确递增时比一次：不同就是计数器有问题。

小结

index_pos 由调用方维护，forward 不核对它和缓存长度。
清了缓存却不归零：多 token 输入时 mask 按旧位置建，报 cannot broadcast [13, 45] to [1, 9, 13, 13]；单 token 时位置整体平移，RoPE 下结果不变。
计数器按调用次数而不是 token 数前进：同一段对话里位置差错了，静默生成乱码。把清缓存和归零写在一个方法里，计数器按 token 数前进。

cannot broadcast [13, 45] to [1, 9, 13, 13]

现象

出处

修法

怎么确认没中招

小结

💬 关联讨论

登录