🐛 [Bug] partitioning ignores device placement, so an engine absorbs host work and returns it on the GPU

[SrivastavaKshitij]

Bug Description

Partitioning never consults device placement. A .to(device="cpu", dtype=...) that a model
used to move a branch to the host is happily absorbed into a TensorRT engine, the device part
of the cast is dropped, and the result comes back on the GPU -- where it meets the
host-resident siblings that were supposed to be its peers.

Three places, none of which looks at device:

1. py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py :: to_copy_dtype_validator
   admits aten._to_copy on dtype alone:

    def validate_dtype(to_copy_node: Node) -> bool:
        """Returns true if the to_copy node can be converted to TRT
   
        Based on data type being casted to
        """
        allowed_casts = {
            torch.float,
            torch.int32,
            torch.int64,
            torch.bool,
            torch.int8,
            torch.float16,
            torch.bfloat16,
        }
   
        # Validate input node has convertible kwargs
        if "dtype" in to_copy_node.kwargs:
            if to_copy_node.kwargs["dtype"] in allowed_casts:
                return True

   to_copy_node.kwargs["device"] is never read, so .to(device="cpu", dtype=torch.int32)
   validates exactly like .to(dtype=torch.int32).

2. py/torch_tensorrt/dynamo/conversion/impl/cast.py :: to_copy has no device parameter at
   all -- its signature is (ctx, target, source_ir, name, input, dtype, force_layer) -- so the
   device half of the cast has nowhere to go and is silently discarded.

3. py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py :: is_node_supported
   weighs converter-registry membership and torch_executed_ops only. In the version tested
   the sole additional capability check is for complex dtypes:

        if TorchTensorRTOperatorSupport._has_complex_dtype(node):
            # Complex-dtype tensors are not supported by TensorRT; force PyTorch fallback

   There is no analogous check for placement, so a node whose recorded meta["val"] lives on
   the CPU is treated as a candidate for an engine like any other.

Everything on that path then returns on the device. The engine's meta kernel builds its fake
outputs on the input's device (runtime/meta_ops/register_meta_ops.py:
torch.empty(output_shape, dtype=info["dtype"], device=inputs[0].device)), and the real engine
writes to CUDA. inline_trt_modules in dynamo/_exporter.py wires the engine result into its
consumers without re-placing it on the device the graph recorded.

The recorded placement is right there in the metadata and is simply not honored. In the run
below the compiled module prints:

_run_on_acc_0 recorded output devices: [device(type='cpu')]

and then hands its consumers a CUDA tensor.

To Reproduce

docker run --rm --gpus all --ipc=host -v "$PWD":/w -w /w \
        nvcr.io/nvidia/pytorch:26.07-py3 python repro.py

repro.py

import sys
import traceback

import torch
import torch_tensorrt
from torch._subclasses.fake_tensor import FakeTensorMode
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch_tensorrt.dynamo._exporter import transform

ROWS = 16


@torch.library.custom_op("repro::host_index", mutates_args=())
def host_index(counts: torch.Tensor) -> torch.Tensor:
    """Returns a host-resident index vector. Has no TensorRT converter, so it stays in Torch."""
    return torch.arange(0, counts.shape[0], device="cpu", dtype=torch.int32)


@host_index.register_fake
def _host_index_fake(counts: torch.Tensor) -> torch.Tensor:
    return torch.empty(counts.shape[0], dtype=torch.int32, device="cpu")


@torch.library.custom_op("repro::host_cat", mutates_args=())
def host_cat(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """Concatenates two host tensors. Has no converter, so the engine boundary lands here."""
    return torch.cat([a, b])


@host_cat.register_fake
def _host_cat_fake(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return torch.cat([a, b])


class HostBranchAbsorbed(torch.nn.Module):
    """Moves a branch to the host with `.to(device="cpu", dtype=...)`, then uses it there."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        counts = (x * 2.0).to(device="cpu", dtype=torch.int32)
        limit = torch.ops.repro.host_index(counts)
        return torch.ops.repro.host_cat(counts, limit)


def build_engine() -> torch.fx.GraphModule:
    """Compiles the model; the host branch is absorbed into the engine."""
    model = HostBranchAbsorbed().eval().cuda()
    x = torch.randn((ROWS,), device="cuda")
    exported = torch.export.export(model, (x,))
    return torch_tensorrt.dynamo.compile(
        exported,
        inputs=(x,),
        min_block_size=1,
        pass_through_build_failures=True,
    )


def _is_device_mismatch(exc: BaseException) -> bool:
    """Returns True if `exc` reports a cpu/cuda placement conflict."""
    msg = str(exc)
    return (
        "cpu" in msg
        and "cuda" in msg
        and ("expected device" in msg or "same device" in msg or "Device Propagation" in msg)
    )


def main(argv: list[str] | tuple[str, ...] = ()) -> int:
    """Runs the compiled engine eagerly and through the inlined export path."""
    del argv
    print(f"torch          {torch.__version__}")
    print(f"torch_tensorrt {torch_tensorrt.__version__}")

    trt_gm = build_engine()
    print(trt_gm.graph, flush=True)
    for name, child in trt_gm.named_children():
        exprs = getattr(child, "symbolic_shape_expressions", None)
        if exprs is not None:
            print(f"{name} inputs : {exprs.get('inputs')}")
            print(f"{name} outputs: {exprs.get('outputs')}", flush=True)
    for node in trt_gm.graph.nodes:
        val = node.meta.get("val")
        if isinstance(val, list):
            print(
                f"{node.name} recorded output devices: {[getattr(v, 'device', None) for v in val]}"
            )

    print("\n===== eager run of the compiled module =====", flush=True)
    eager_failed = False
    try:
        out = trt_gm(torch.randn((ROWS,), device="cuda"))
        print(f"eager returned {out.device} {tuple(out.shape)}")
    except Exception as exc:  # pylint: disable=broad-except
        traceback.print_exc()
        eager_failed = _is_device_mismatch(exc)

    print("\n===== inlined (export) run under FakeTensorMode =====", flush=True)
    inlined = transform(trt_gm)
    inlined.recompile()
    print(inlined.graph, flush=True)
    fake_failed = False
    with FakeTensorMode(shape_env=ShapeEnv()):
        x_fake = torch.empty((ROWS,), dtype=torch.float32, device="cuda")
        try:
            out = inlined(x_fake)
            print(f"inlined graph returned {out}")
        except Exception as exc:  # pylint: disable=broad-except
            traceback.print_exc()
            fake_failed = _is_device_mismatch(exc)

    reproduced = eager_failed or fake_failed
    print(f"\neager device mismatch:  {eager_failed}")
    print(f"inlined device mismatch: {fake_failed}")
    print(f"\nreproduced: {reproduced}")
    return 0 if reproduced else 1


if __name__ == "__main__":
    sys.exit(main(argv=sys.argv))

output

torch          2.13.0a0+9186a08b2c.nv26.07
torch_tensorrt 2.14.0a0
graph():
    %x : [num_users=1] = placeholder[target=x]
    %_run_on_acc_0 : [num_users=1] = call_module[target=_run_on_acc_0](args = (%x,), kwargs = {})
    %_run_on_gpu_1 : [num_users=1] = call_module[target=_run_on_gpu_1](args = (%_run_on_acc_0,), kwargs = {})
    return (_run_on_gpu_1,)
_run_on_acc_0 inputs : [{'shape_exprs': [16], 'dtype': torch.float32, 'name': 'x'}]
_run_on_acc_0 outputs: [{'shape_exprs': [16], 'dtype': torch.int32}]
_run_on_acc_0 recorded output devices: [device(type='cpu')]

===== eager run of the compiled module =====
Traceback (most recent call last):
  File "/w/repro.py", line 132, in main
    out = trt_gm(torch.randn((ROWS,), device="cuda"))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 1000, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[...]
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<eval_with_key>.30", line 7, in forward
    _run_on_gpu_1 = self._run_on_gpu_1(_run_on_acc_0);  _run_on_acc_0 = None
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 1000, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[...]
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<eval_with_key>.32", line 6, in forward
    host_cat = torch.ops.repro.host_cat.default(_to_copy, host_index);  _to_copy = host_index = None
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 875, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
[...]
  File "/usr/local/lib/python3.12/dist-packages/torch/_library/custom_ops.py", line 502, in wrapped_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/w/repro.py", line 69, in host_cat
    return torch.cat([a, b])
           ^^^^^^^^^^^^^^^^^
RuntimeError: Expected all tensors to be on the same device, but got tensors is on cpu, different from other tensors on cuda:0 (when checking argument in method wrapper_CUDA_cat)

===== inlined (export) run under FakeTensorMode =====
WARNING:py.warnings:/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/_exporter.py:500: UserWarning: Attempted to insert a get_attr Node with no underlying reference in the owning GraphModule! Call GraphModule.add_submodule to add the necessary submodule, GraphModule.add_parameter to add the necessary Parameter, or nn.Module.register_buffer to add the necessary buffer
  engine_node = gm.graph.get_attr(engine_name)

graph():
    %x : [num_users=1] = placeholder[target=x]
    %_run_on_acc_0_engine : [num_users=1] = get_attr[target=_run_on_acc_0_engine]
    %execute_engine_default : [num_users=1] = call_function[target=torch.ops.tensorrt.execute_engine.default](args = ((%x,), %_run_on_acc_0_engine), kwargs = {})
    %getitem : [num_users=2] = call_function[target=operator.getitem](args = (%execute_engine_default, 0), kwargs = {})
    %host_index : [num_users=1] = call_function[target=torch.ops.repro.host_index.default](args = (%getitem,), kwargs = {})
    %host_cat : [num_users=1] = call_function[target=torch.ops.repro.host_cat.default](args = (%getitem, %host_index), kwargs = {})
    return (host_cat,)
Traceback (most recent call last):
  File "/w/repro.py", line 146, in main
    out = inlined(x_fake)
          ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 1000, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[...]
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<eval_with_key>.36", line 10, in forward
    host_cat = torch.ops.repro.host_cat.default(getitem, host_index);  getitem = host_index = None
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 875, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
[...]
  File "/usr/local/lib/python3.12/dist-packages/torch/_library/custom_ops.py", line 784, in fake_impl
    return self._abstract_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/w/repro.py", line 74, in _host_cat_fake
    return torch.cat([a, b])
           ^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_compile.py", line 54, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[...]
  File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1097, in merge_devices
    raise FakeTensorDeviceMismatchError(func, common_device, t.device)
torch._subclasses.fake_tensor.FakeTensorDeviceMismatchError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

eager device mismatch:  True
inlined device mismatch: True

reproduced: True

Expected behavior

Device placement should be part of the partitioning decision, and a cast that changes device
should not be claimed by a converter that cannot perform the change:

• to_copy_dtype_validator should decline aten._to_copy when kwargs["device"] names a
  device other than the engine's, instead of validating on dtype alone. Declining is the
  right outcome on the merits -- a TensorRT engine cannot produce a host tensor -- and it is
  the outcome the partitioner is designed to handle: the node falls back to Torch and the rest
  of the model still compiles.
• More generally, is_node_supported should treat a recorded placement that differs from the
  engine's device the way it already treats complex dtypes: force the fallback rather than
  absorb the node.
• Wherever an engine output's recorded meta["val"] device differs from the device the engine
  writes to, the boundary needs an explicit transfer -- and the meta kernel needs to report the
  same device, rather than deriving it from inputs[0].device.

Separately, and as a second request: the diagnostics here are poor. Compilation reports success,
the recorded metadata says cpu, and the first sign of trouble is a device mismatch inside an
unrelated consumer. Silently dropping the device argument of a cast should at minimum produce
a warning naming the node.

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• Pytorch NGC container : 26.07-py3