RuntimeError: Expected all tensors to be on the same device, but got weight is on cpu

快速结论：这个报错通常出现在 MiniMax-H3 视频 VAE 解码阶段，当 VAE 处于部分 offload 状态、qk_norm_scale 停留在 CPU 而上游张量在 CUDA/ROCm 设备上时触发。优先排查 VAE 是否被部分卸载，以及 qk_norm_scale 与 query 的设备是否一致。

适用环境：已在以下环境复现并验证：ComfyUI v0.37.0（73c9bad4）、comfy-kitchen 0.2.35、AMD Radeon R9700（gfx1201，32 GB VRAM）、ROCm 7.2、torch 2.12.1+rocm7.2、Linux；使用 --reserve-vram 1 --fp32-vae，测试 INT8 attention 路径时另加 --use-ck-attention。

最快修复方案：在调用 rms_rope 之前，将 self.qk_norm_scale 按需移动到 query.device，Issue 中已在该环境下验证可完成解码（示例补丁见“解决步骤”）。

注意事项：该修复只是让调用点不再传入 CPU 张量，comfy_kitchen eager 后端本身仍然不做设备对齐，其他走 eager fallback 的路径仍可能遇到同类问题。补丁行为为本地验证，非官方确认方案。

问题场景

在 ComfyUI 中运行 MiniMax-H3 Ref2VA（pruned INT8 convrot）工作流，加载 minimax_h3_video_vae_fp16.safetensors。前置的条件、采样和音频 VAE 加载均正常，崩溃发生在视频 VAE 解码阶段。触发的显卡为 32 GB 显存的 AMD Radeon R9700，VAE 无法完全驻留显存。

调用链为：comfy/ldm/minimax/vae.py 的 forward → rms_rope → comfy/quant_ops.py 中的 rms_rope_split_half_ → torch.ops.comfy_kitchen.rms_rope_split_half_ → comfy_kitchen/backends/eager/rope.py 的 _rms_rope1 → torch.nn.functional.rms_norm。

报错原文

RuntimeError: Expected all tensors to be on the same device, but got weight is on cpu,
different from other tensors on cuda:0 (when checking argument in method wrapper_CUDA___fused_rms_norm)

原因分析

最可能的原因是 comfy-kitchen 的后端路由与 eager 后端缺少设备对齐共同导致的。

后端优先级为 registry.set_priority(["hip", "cuda", "triton", "eager"])，而 rms_rope_split_half_ 对各后端声明的 default_devices 如下：hip 接受 {'cuda', 'hip'}，cuda 接受 {'cuda'}，triton 接受 {'cuda', 'xpu'}，eager 接受 {'cpu', 'mps', 'cuda', '*', 'meta', 'hpu', 'xpu'}。

由于 qk_norm_scale 位于 CPU，hip/cuda/triton 后端均无法接管该调用，调用会落到唯一接受 CPU 的 eager 后端；而 eager 实现直接调用 F.rms_norm(x, weight=scale)，从未把 scale 移动到 x.device，于是出现设备不一致报错。

同文件中 _fused_norm_pad 通过 comfy.ops.cast_bias_weight(norm, x, offloadable=True) 逐次调用做转换，而 rms_rope 调用传入的是原始 buffer，这正是差异所在。

环境排查

确认 ComfyUI 版本，Issue 环境为 v0.37.0（73c9bad4），调用点与 master 一致。
确认 comfy-kitchen 版本，Issue 环境为 0.2.35。
确认 ROCm 与 PyTorch 版本，Issue 环境为 ROCm 7.2 / torch 2.12.1+rocm7.2。
确认显卡与显存，Issue 环境为 AMD Radeon R9700（gfx1201），32 GB VRAM。
确认操作系统，Issue 环境为 Linux（workstation）。
确认启动参数，Issue 环境使用 --reserve-vram 1 --fp32-vae，测试 INT8 attention 路径时使用 --use-ck-attention。
确认 VAE 是否处于部分驻留/部分 offload 状态（同模型在生产构建日志中出现 Unloaded partially: ... MB freed, ... MB buffer reserved）。
确认触发模型，Issue 中为 MiniMax-H3 Ref2VA（pruned INT8 convrot）与 VAE minimax_h3_video_vae_fp16.safetensors。

解决步骤

确认问题可稳定复现，且报错发生在视频 VAE 解码阶段而非加载或采样阶段。
按 Issue 中验证的补丁，在 comfy/ldm/minimax/vae.py 调用 rms_rope 之前，将 qk_norm_scale 按需移动到 query.device：

_qk_scale = self.qk_norm_scale
if _qk_scale.device != query.device:
    _qk_scale = _qk_scale.to(query.device)
query, key = rms_rope(
    query, key, rotary_pos_emb, _qk_scale,
    ...)

如果只是需要最小复现来确认设备不匹配问题，可在不加载模型的情况下运行 Issue 中的片段：scale_cpu 传入会报 weight is on cpu，scale_gpu 传入则正常。可优先尝试用这个片段确认 eager 后端对 CPU 张量的行为。
如果无法修改代码，可优先尝试让 VAE 完整驻留显存（即避免部分 offload），以减少 qk_norm_scale 落在 CPU 的概率；但 Issue 未验证该绕过方式的具体配置。

验证方法

使用相同工作流在相同环境下重跑（Issue 中为 960×544、192 帧、8-step dual-clock、相同 seed）。修复前在 VAE 解码阶段抛出 RuntimeError: ... weight is on cpu ... (wrapper_CUDA___fused_rms_norm)；修复后完整执行结束，日志显示 Prompt executed in 00:15:53，输出 192 帧 @24fps 且含音轨，全程不再出现 weight is on cpu。