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

快速结论：这个报错通常出现在 ComfyUI 运行 MiniMax-H3 视频 VAE decode、且 VAE 未被完整装载到显存（部分 offload）的场景。优先排查 qk_norm_scale 这个 buffer 是否停留在 CPU，以及当前 rms_rope_split_half_ 是否因 CPU 参数回退到了 comfy-kitchen 的 eager 后端。

适用环境：Issue 已确认环境为 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 workstation；相关启动参数为 --reserve-vram 1 --fp32-vae，测试 INT8 attention 路径时另加 --use-ck-attention。复现脚本使用极小张量，Issue 说明在 CUDA 上同样可复现。

最快修复方案：Issue 中已在本地验证：在 rms_rope 调用处先把 self.qk_norm_scale 搬到 query.device，再传给 rms_rope。这是 call-site 层面的 workaround，不是对 eager 后端本身的修复。

注意事项：该补丁只解决当前解码路径不再走错设备的问题；comfy-kitchen 的 eager fallback 本身仍然不做 device alignment，其他调用方仍可能踩到同一底层缺口。修改 ComfyUI 源码会在升级时被覆盖，需要自行记录或维护补丁。

问题场景

用户在 ComfyUI 中运行 MiniMax-H3 Ref2VA（pruned INT8 convrot）工作流，VAE 使用 minimax_h3_video_vae_fp16.safetensors。条件、采样和音频 VAE 加载均正常，崩溃发生在 video VAE decode 阶段：comfy/ldm/minimax/vae.py 调用 rms_rope，进入 comfy_kitchen 的 rms_rope_split_half_。

触发前提是 VAE 在 32 GB 显卡上只能部分驻留（日志中出现 Unloaded partially: ... MB freed, ... MB buffer reserved），导致 qk_norm_scale 这个注册 buffer 到达 kernel 时仍在 CPU，而 query 已经在 cuda:0。

报错原文

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)

调用栈关键位置：

File ".../comfy/ldm/minimax/vae.py", line 301, in forward
    query, key = rms_rope(
File ".../comfy/quant_ops.py", in rms_rope_split_half_
    torch.ops.comfy_kitchen.rms_rope_split_half_(...)
File ".../comfy_kitchen/backends/eager/rope.py", line 85, in _rms_rope1
    x_norm = torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=scale, eps=epsilon)

原因分析

根因是 backend routing 与 eager fallback 缺少设备对齐，而不是调用点本身写错。

rms_rope_split_half_ 的后端优先级为 registry.set_priority(["hip", "cuda", "triton", "eager"])，各后端声明的 default_devices 为：hip 接受 {'cuda', 'hip'}，cuda 接受 {'cuda'}，triton 接受 {'cuda', 'xpu'}，eager 接受 {'cpu', 'mps', 'cuda', '*', 'meta', 'hpu', 'xpu'}。当参数中有 CPU 张量时，hip/cuda/triton 都无法接手，调用会落到唯一接受 CPU 的 eager 后端。

eager 后端的 _rms_rope1 直接调用 F.rms_norm(x, weight=scale)，没有把 scale 移到 x.device，于是 HIP 后端在参数落在 CPU 时拒绝调用并抛出上述错误。

同一文件中的 _fused_norm_pad 已经通过 comfy.ops.cast_bias_weight(norm, x, offloadable=True) 按调用做转换，但 rms_rope 这里传的是原始 buffer。因此这是“部分 offload VAE + CPU-resident buffer + eager fallback 不做 alignment”共同导致的崩溃。

环境排查

确认 ComfyUI 版本：Issue 为 v0.37.0（73c9bad4），并注明调用点在 master 上相同。
确认 comfy-kitchen 版本：0.2.35。
确认 GPU 与后端：AMD Radeon R9700（gfx1201）、ROCm 7.2。
确认 PyTorch：torch 2.12.1+rocm7.2。
确认操作系统：Linux workstation。
确认启动参数：--reserve-vram 1 --fp32-vae；测试 INT8 attention 路径时包含 --use-ck-attention。
确认模型与 VAE：MiniMax-H3 Ref2VA（pruned INT8 convrot），VAE 为 minimax_h3_video_vae_fp16.safetensors。
确认 VAE 是否部分驻留：32 GB 显卡上出现 Unloaded partially: ... MB freed, ... MB buffer reserved 日志时，更容易触发。

解决步骤

定位 comfy/ldm/minimax/vae.py 中 rms_rope 的调用位置（Issue 为第 301 行附近）。
在调用前取出 self.qk_norm_scale，判断其 device 是否与 query.device 一致。
若不一致，先把 scale 转换到 query.device，再把转换后的张量传给 rms_rope。Issue 中验证过的补丁形式如下：

             _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 中被标记为 Verified locally，属于调用点 workaround，可优先尝试。底层 eager fallback 的设备不对齐问题仍然存在。

验证方法

用与崩溃时相同的工作流复测（Issue 中的对照条件为 960×544、192 frames、8-step dual-clock、相同 seed）：

修复前：VAE decode 阶段出现 RuntimeError: ... weight is on cpu ... (wrapper_CUDA___fused_rms_norm)。
修复后：流程完成，Issue 记录为 Prompt executed in 00:15:53，192 frames @24fps，音频轨正常，日志中不再出现 weight is on cpu。

也可以用 Issue 提供的最小复现脚本单独验证：scale_cpu 会触发报错，scale_gpu 正常通过；补丁生效后，同一路径不应再把 CPU 张量交给需要 CUDA 的算子。