On cu130: Can't export tensors on a different CUDA device index. cu128 works in light testing

[john815308]

A flux 2 dev workflow when I put the text model (mistral) on the cuda 1 device on the ClipLoaderMultiGpu, that node works but then fails on the Clip Text Encode saying "Can't export tensors on a different CUDA device index" Using cuda 12.8 this works in light testing. Using 130 it always fails as described.

[fAIseh00d]

Same issue, updated from torch2.4 cuda12.8 to torch 2.9.1 cuda13.0 and the issue appeared.

Workflow:
z-image-test.json

Log

got prompt
Requested to load Lumina2
loaded completely; 21949.99 MB usable, 11739.54 MB loaded, full load: True
0%|                                                         | 0/9 [00:00<?, ?it/s]
!!! Exception during processing !!! Can't export tensors on a different CUDA device index. Expected: 1. Current device: 0.
Traceback (most recent call last):
File "/mnt/nvme/ComfyUI/execution.py", line 518, in execute
  output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
                                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/execution.py", line 329, in get_output_data
  return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/execution.py", line 303, in _async_map_node_over_list
  await process_inputs(input_dict, i)
File "/mnt/nvme/ComfyUI/execution.py", line 291, in process_inputs
  result = f(**inputs)
           ^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/nodes.py", line 1544, in sample
  return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/nodes.py", line 1511, in common_ksampler
  samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/sample.py", line 60, in sample
  samples = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 1178, in sample
  return sample(self.model, noise, positive, negative, cfg, self.device, sampler, sigmas, self.model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 1068, in sample
  return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 1050, in sample
  output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=latent_shapes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/patcher_extension.py", line 112, in execute
  return self.original(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 994, in outer_sample
  output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=latent_shapes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 980, in inner_sample
  samples = executor.execute(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/patcher_extension.py", line 112, in execute
  return self.original(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 752, in sample
  samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
  return func(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/k_diffusion/sampling.py", line 1522, in sample_er_sde
  denoised = model(x, sigmas[i] * s_in, **extra_args)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 401, in __call__
  out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 953, in __call__
  return self.outer_predict_noise(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 960, in outer_predict_noise
  ).execute(x, timestep, model_options, seed)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/patcher_extension.py", line 112, in execute
  return self.original(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 963, in predict_noise
  return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 381, in sampling_function
  out = calc_cond_batch(model, conds, x, timestep, model_options)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 206, in calc_cond_batch
  return _calc_cond_batch_outer(model, conds, x_in, timestep, model_options)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 214, in _calc_cond_batch_outer
  return executor.execute(model, conds, x_in, timestep, model_options)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/patcher_extension.py", line 112, in execute
  return self.original(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/samplers.py", line 326, in _calc_cond_batch
  output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/model_base.py", line 163, in apply_model
  return comfy.patcher_extension.WrapperExecutor.new_class_executor(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/patcher_extension.py", line 112, in execute
  return self.original(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/model_base.py", line 205, in _apply_model
  model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
  return self._call_impl(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
  return forward_call(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/ldm/lumina/model.py", line 600, in forward
  return comfy.patcher_extension.WrapperExecutor.new_class_executor(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/patcher_extension.py", line 112, in execute
  return self.original(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/ldm/lumina/model.py", line 635, in _forward
  img, mask, img_size, cap_size, freqs_cis = self.patchify_and_embed(x, cap_feats, cap_mask, adaln_input, num_tokens, transformer_options=transformer_options)
                                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/ldm/lumina/model.py", line 581, in patchify_and_embed
  cap_feats = layer(cap_feats, cap_mask, freqs_cis[:, :cap_pos_ids.shape[1]], transformer_options=transformer_options)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
  return self._call_impl(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
  return forward_call(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/ldm/lumina/model.py", line 295, in forward
  clamp_fp16(self.attention(
             ^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
  return self._call_impl(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
  return forward_call(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/ComfyUI/comfy/ldm/lumina/model.py", line 114, in forward
  xq, xk = apply_rope(xq, xk, freqs_cis)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/comfy_kitchen/__init__.py", line 179, in apply_rope
  return torch.ops.comfy_kitchen.apply_rope(xq, xk, freqs_cis)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/_ops.py", line 1255, in __call__
  return self._op(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/_library/custom_ops.py", line 343, in backend_impl
  result = self._backend_fns[device_type](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/_compile.py", line 53, in inner
  return disable_fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 1044, in _fn
  return fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/_library/custom_ops.py", line 376, in wrapped_fn
  return fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/comfy_kitchen/backends/eager/rope.py", line 33, in _op_apply_rope
  return impl(**kwargs)
         ^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/comfy_kitchen/backends/cuda/__init__.py", line 439, in apply_rope
  _wrap_for_dlpack(xq),
  ^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/comfy_kitchen/backends/cuda/__init__.py", line 135, in _wrap_for_dlpack
  return tensor.__dlpack__(stream=-1)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/nvme/venv/lib/python3.12/site-packages/torch/_tensor.py", line 1726, in __dlpack__
  raise BufferError(
BufferError: Can't export tensors on a different CUDA device index. Expected: 1. Current device: 0.

Prompt executed in 3.11 seconds

[joecomputerjoe2]

comfyui-multigpu-fixed-cu130-dlpack.zip

i had chatgpt analyze the code and propose a fix, it took a few tries and a couple small adjustments to get it working but it works!
remove the custom node folder and replace it with the one from the zip, tested with qwen image to offload the clip from my 3090 to my second video card (3060) im sure the maintainers will add an update for this eventually but wanted to add this here for anyone like me who got used to using this for their workflows and wants to make it work again asap

-joe

[joecomputerjoe2]

p.s. can confirm the issue was caused by upgrading to cu130

[fantasyz]

@joecomputerjoe2 It works. Thank you. 👍
I diffed your files with mine. It seems that only the __init__.py is updated.

Only this particular piece of code inserted at line 31:

# --- cu130 / PyTorch DLPack device-guard patch ---
# comfy_kitchen's CUDA backend uses tensor.__dlpack__ which requires the current CUDA device
# to match the tensor's device index. With multi-GPU setups the global current device can differ,
# causing: BufferError: Can't export tensors on a different CUDA device index.
def _patch_comfy_kitchen_dlpack_device_guard():
    try:
        import comfy_kitchen.backends.cuda as ck_cuda  # type: ignore
    except Exception:
        return  # comfy_kitchen not installed or no CUDA backend

    try:
        orig_wrap = getattr(ck_cuda, "_wrap_for_dlpack", None)
        if orig_wrap is None or getattr(orig_wrap, "_multigpu_patched", False):
            return

        def _wrap_for_dlpack_guarded(tensor):
            try:
                if torch.cuda.is_available() and hasattr(tensor, "is_cuda") and tensor.is_cuda:
                    idx = tensor.device.index
                    if idx is not None:
                        # Ensure current CUDA device matches the tensor before exporting via DLPack.
                        torch.cuda.set_device(idx)
            except Exception:
                pass
            return orig_wrap(tensor)

        _wrap_for_dlpack_guarded._multigpu_patched = True  # type: ignore
        setattr(ck_cuda, "_wrap_for_dlpack", _wrap_for_dlpack_guarded)
        logger.info("[MultiGPU] Patched comfy_kitchen CUDA DLPack wrapper with device guard.")
    except Exception as e:
        logger.debug(f"[MultiGPU] Failed to patch comfy_kitchen DLPack wrapper: {e}")

_patch_comfy_kitchen_dlpack_device_guard()
# --- end patch ---