RuntimeError: Tensor device mismatch … found at least two devices, neuron:0 and cpu! (scaled_dot_product_attention)

快速结论：这个报错通常出现在 Transformers 的 generate() 使用 cache_implementation="static" 时：如果输入张量留在 CPU（host-side），而模型已在加速设备（如 neuron:0）上，嵌套在 dict 里的 attention mask 张量不会被搬到模型设备，最终在 SDPA 处触发设备不一致。优先排查输入张量是否全部在模型设备上，以及 config 是否带 layer_types（会返回按层类型分组的 mask dict）。

适用环境：Issue 原文确认环境为 transformers 5.18.0（main 上未变）、torch 2.14.0 + torch_neuronx 2.14.3、Python 3.14、AWS trn2。另有一位复现者用 RTX 5090、Python 3.12.13、PyTorch 2.11.0+cu130 独立确认了同一类 CPU/GPU 设备不匹配，并使用小型随机初始化的 Gemma3ForCausalLM。

最快修复方案：暂无确认的一步修复方案。Issue 作者表示已自行准备两个 bug 的修复并会提交 PR，但尚未在该讨论链中给出可直接套用的补丁或版本。

注意事项：Issue 中的验证在 Neuron 和 GPU 上分别只覆盖了部分场景（GPU 复现时关闭了编译，未测试 Neuron 与 TP）。把嵌套 dict 张量手动搬到模型设备可作为临时绕过思路，但属于推断而非 Issue 中已验证的方案；修改 transformers 内部 prepare_inputs_for_generation / static cache 初始化形状相关逻辑前请自行确认版本一致性。

问题场景

用户在调用 Transformers 的 generate() 并指定 cache_implementation="static" 时触发问题。复现方式是从 CPU 侧构造输入（tokenizer 返回的 inputs 留在 host 上），模型已 .to("neuron")，然后直接 model.generate(**inputs, ...)。在带 layer_types 的 config（多 pattern 配置如 Gemma 3、Qwen3.5，以及恢复 dict 返回后的单 pattern Qwen3）下，可编译 cache 的 attention mask 是按层类型分组的 dict，dict 内部的张量不会被搬到模型设备。此外，Issue 还报告了另一类问题：在 TP 下设置 prefill_chunk_size 时，早期 static cache 的 KV heads 数量被错误除以 tp_size。

报错原文

RuntimeError: Tensor device mismatch ... found at least two devices, neuron:0 and cpu! (scaled_dot_product_attention)

TP 场景下的另一条报错：

ValueError: Incompatible shapes for broadcasting: (1, 8, 512, 128) and requested shape (1, 4, 512, 128)

另有复现者在 GPU 上看到：

attn_bias is on cpu, different from other tensors on cuda:0

原因分析

原因一（与上面的 device mismatch 报错直接相关，也是 GPU 复现者确认的部分）：prepare_inputs_for_generation 的第 6 步只把 model_inputs 顶层的张量搬到模型设备，嵌套在 dict 里的张量不会被处理。使用可编译 cache 时，attention mask 往往是一个以层类型为 key 的 dict（config 含 layer_types 时），因此 dict 内部的 mask 张量仍停留在 CPU。只要输入来自 host 且 batch 有 padding（mask 不会被跳过），SDPA 就会遇到 CPU 与 neuron:0（或 cuda:0）两个设备，抛出上述错误。该问题在多 pattern 配置中自 #46590 起存在，在单 pattern 配置中自 #47019 恢复 dict 返回后存在。

原因二（对应广播形状报错，属于 TP 场景的独立 bug）：设置 prefill_chunk_size 时，_get_static_cache_init_shape 会把 KV heads 除以 tp_size，即使 plan 并不对 attention 做切分（没有 plan，或使用 colwise_gather_output）。期望行为是仅在 k_proj 以 colwise 方式被切分时才除以 KV heads。Issue 还提到 #47811 会把这种切分风格改名为 colwise_units（在复制的 KV heads 场景下）。

环境排查

确认 transformers 版本：Issue 为 5.18.0，并声明在 main 上未变。
确认 torch 与设备后端：torch 2.14.0 + torch_neuronx 2.14.3；GPU 复现侧为 PyTorch 2.11.0+cu130。
确认 Python 版本：Issue 为 3.14，GPU 复现侧为 3.12.13。
确认硬件：AWS trn2；GPU 复现侧为 RTX 5090。
确认是否使用 cache_implementation="static"，以及是否设置了 prefill_chunk_size。
确认模型 config 是否带 layer_types（影响 attention mask 是否为 dict）。
确认输入张量所在设备：tokenizer 输出是否留在 CPU，模型是否已在 neuron/cuda 上。
确认 batch 是否为 padded batch（影响 mask 是否被跳过）。
TP 场景下确认 tp_size、是否有 attention plan，以及 k_proj 是否以 colwise 方式被切分。
GPU 复现者还确认：相同模型使用 on-device 输入时通过，dynamic cache 下两种放置方式都通过，tiny Llama 对照在两种 cache 下均通过。

解决步骤

先按 Issue 最小复现路径确认触发条件：从 CPU 侧构造输入，模型放在 neuron 上，使用 cache_implementation="static" 调 generate()。如果使用 GPU，可参照 GPU 复现者的做法用小型随机初始化 Gemma3ForCausalLM，并关闭编译做正确性验证。
检查报错模型配置中是否包含 layer_types，从而判断 attention mask 是否为按层类型分组的 dict。
对照 GPU 复现者的结论做交叉验证：同一模型改用 on-device 输入，或改用 dynamic cache，观察是否不再报 device mismatch。这两种情况通过、仅 host 输入 + static cache 失败，基本可定位到嵌套 dict 张量未搬运的问题。
可优先尝试的临时绕过（非 Issue 验证过的正式修复）：在进入 generate() / forward 前，把自身构造的输入张量以及嵌套 dict 中的张量都显式搬到模型设备，再调用生成；Issue 作者指出“把 layer-mask dict 的张量随 forward 输入一起搬运可以修复该失败，并得到完全一致的生成 token ID”，但这是复现者的验证描述，不是官方已合并的补丁。
若报错是 TP 下的广播形状错误，检查是否设置了 prefill_chunk_size，并确认 plan 是否真的对 attention 做了切分；Issue 期望的行为是只在 k_proj 以 colwise 方式切分时才把 KV heads 除以 tp_size。
关注 Issue 作者提到的两个修复 PR；在官方修复落地前不要直接改动本地 site-packages 中的 layout 相关代码，除非你能对照具体版本核对 prepare_inputs_for_generation、_get_static_cache_init_shape 的实现。

验证方法

用同一份最小复现脚本对比三种情况下是否通过：host-side 输入 + cache_implementation="static"（应触发报错）、on-device 输入 + static cache（GPU 复现者报告通过）、以及 host-side 输入 + dynamic cache（GPU 复现者报告通过）。如果应用临时搬运方案后，生成不再出现 Tensor device mismatch ... neuron:0 and cpu，且生成的 token ID 与 on-device 输入版本完全一致，说明问题被绕过。TP 场景下则检查 broadcast 形状报错是否消失，并确认 cache 初始化后的 KV heads 数量与 plan 是否切分 attention 一致。