Host-resident TensorRT shape-tensor inputs incur a device round trip and a per-call stream sync

[SrivastavaKshitij]

Bug Description

A TensorRT shape tensor input is read from host memory, but setup_input_tensors in
core/runtime/execute_engine.cpp requires every input to be on the device and then copies a
shape input back to the host with .cpu():

TORCHTRT_CHECK(
    inputs[i].is_cuda(), "Expected input tensors to have device cuda, found device " << inputs[i].device());
...
if (binding.is_shape_tensor) {
  auto input_cpu = inputs[i].clone().contiguous().cpu().to(torch::kInt64);

execute_engine also moves any non-CUDA input to the device up front, so a shape value the
caller already holds on the host is pushed device-side and then pulled straight back. The .cpu()
is a stream synchronization: the host blocks until every kernel already queued on the stream has
finished before the engine is enqueued. When the shape value is produced on the host (e.g. an
AOTInductor wrapper that materializes a symbolic dimension with scalar_tensor), the whole round
trip is wasted, and the sync serializes host and device on a path that has no reason to.

To Reproduce

nvcr.io/nvidia/pytorch:26.07-py3, one L4. The script builds a one-engine model whose engine
takes a TensorRT shape-tensor input (arange(0, n, 1) makes n a shape binding) and measures
per-call host overhead with the shape value on the device vs on the host, plus the isolated cost
of the copy on a stream that has work queued on it.

benchmark_host_shape_inputs.py

"""Benchmark for host shape-tensor inputs to a TensorRT engine.

A TensorRT *shape tensor* input is read from host memory. ``setup_input_tensors`` in
``core/runtime/execute_engine.cpp`` nonetheless requires every input to be on the device
and copies a shape input back to the host with ``.cpu()``. When the caller already holds
the value on the host, that is an H2D move followed by a D2H copy, and the D2H copy
synchronizes the CUDA stream: the host blocks until the work already queued on the stream
has finished.

This script builds a one-engine model whose engine takes a TensorRT shape-tensor input and
reports two things:

  Part 1 -- per-call host overhead on an idle stream, with the shape input placed on the
            device and on the host. This is the cost the copy adds to every call.
  Part 2 -- the isolated cost of the exact copy the patch changes, measured on a stream that
            has a backlog of work queued on it, to show that the copy waits for that work.

On a runtime that copies a shape input to the host unconditionally, a host-resident shape
input is the slower of the two in Part 1 (an extra H2D then D2H). On a runtime that accepts a
host shape input as is, it is the faster, and Part 2's host path does not wait for the backlog.

Run under ``PYTHONPATH`` pointing at the torch_tensorrt build to compare.
"""

import gc
import statistics
import time
from typing import List, Tuple

import torch
import torch.nn as nn
import tensorrt as trt
import torch_tensorrt
from torch_tensorrt.dynamo.runtime import TorchTensorRTModule

SHAPE_VAL = 64
ITERS = 300
WARMUP = 30
BACKLOG_MATMULS = 20
BACKLOG_SIZE = 4096


class ShapeInputModel(nn.Module):
    """``arange(0, n, 1)`` makes ``n`` a TensorRT shape-tensor input binding."""

    def forward(self, x: torch.Tensor, n: torch.Tensor) -> torch.Tensor:
        a = torch.arange(0, n, 1, device=x.device).to(torch.float32)
        return (x * 2.0 + a).sum(dim=0, keepdim=True)


def build_engine_module() -> TorchTensorRTModule:
    x = torch.randn(SHAPE_VAL, device="cuda")
    n = torch.tensor(SHAPE_VAL, dtype=torch.int64, device="cuda")
    dim = torch.export.Dim("d", min=2, max=256)
    ep = torch.export.export(
        ShapeInputModel().eval().cuda(), (x, n), dynamic_shapes={"x": {0: dim}, "n": None}
    )
    gm = torch_tensorrt.dynamo.compile(
        ep, inputs=[x, n], min_block_size=1, pass_through_build_failures=True,
        use_python_runtime=False, assume_dynamic_shape_support=True,
    )
    m = next(mm for _, mm in gm.named_children() if isinstance(mm, TorchTensorRTModule))
    _assert_shape_binding(m)
    return m


def _assert_shape_binding(m: TorchTensorRTModule) -> None:
    rt = trt.Runtime(trt.Logger(trt.Logger.ERROR))
    ce = rt.deserialize_cuda_engine(m.serialized_engine)
    shape_ins = [
        ce.get_tensor_name(i)
        for i in range(ce.num_io_tensors)
        if ce.get_tensor_mode(ce.get_tensor_name(i)) == trt.TensorIOMode.INPUT
        and ce.is_shape_inference_io(ce.get_tensor_name(i))
    ]
    assert shape_ins, "engine has no shape-tensor input binding"


def ordered_inputs(m: TorchTensorRTModule, x: torch.Tensor, device: str) -> List[torch.Tensor]:
    args = {"x": x, "_local_scalar_dense": torch.tensor(SHAPE_VAL, dtype=torch.int64, device=device)}
    return [args[b] for b in m.input_binding_names]


def per_call_overhead(m: TorchTensorRTModule, x: torch.Tensor, device: str) -> Tuple[float, float]:
    inp = ordered_inputs(m, x, device)
    for _ in range(WARMUP):
        out = torch.ops.tensorrt.execute_engine(inp, m.engine)
    torch.cuda.synchronize()
    ts = []
    gc.disable()  # the call allocates its output; periodic GC otherwise skews the tail
    for _ in range(ITERS):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        torch.ops.tensorrt.execute_engine(inp, m.engine)
        ts.append((time.perf_counter() - t0) * 1e6)
    gc.enable()
    return min(ts), float(out[0])


def copy_cost_under_backlog() -> Tuple[float, float]:
    """Times the two shape-input copies the patch chooses between, on a busy stream."""
    a = torch.randn(BACKLOG_SIZE, BACKLOG_SIZE, device="cuda")
    b = torch.randn(BACKLOG_SIZE, BACKLOG_SIZE, device="cuda")
    c = torch.empty(BACKLOG_SIZE, BACKLOG_SIZE, device="cuda")
    dev_in = torch.tensor(SHAPE_VAL, dtype=torch.int64, device="cuda")
    host_in = torch.tensor(SHAPE_VAL, dtype=torch.int64, device="cpu")

    def queue_backlog() -> None:
        for _ in range(BACKLOG_MATMULS):
            torch.matmul(a, b, out=c)

    def bench(fn) -> float:
        for _ in range(5):
            fn()
        torch.cuda.synchronize()
        ts = []
        for _ in range(50):
            queue_backlog()
            t0 = time.perf_counter()
            fn()
            ts.append((time.perf_counter() - t0) * 1e3)
            torch.cuda.synchronize()
        return statistics.median(ts)

    dev_ms = bench(lambda: dev_in.clone().contiguous().cpu().to(torch.int64))
    host_ms = bench(lambda: host_in.contiguous().to(torch.int64))
    return dev_ms, host_ms


def main() -> None:
    print(f"torch {torch.__version__} | torch_tensorrt {torch_tensorrt.__file__}\n")
    m = build_engine_module()

    x = torch.randn(SHAPE_VAL, device="cuda")
    print(f"Part 1: per-call host overhead, idle stream (min of {ITERS} calls)")
    print(f"{'shape input on':>16} | {'host time':>11} | {'output':>10}")
    print("-" * 44)
    vals = []
    for device in ("cuda", "cpu"):
        mn, val = per_call_overhead(m, x, device)
        vals.append(val)
        print(f"{device:>16} | {mn:>8.1f} us | {val:>10.3f}")
    assert abs(vals[0] - vals[1]) < 1e-3, f"device and host inputs disagree: {vals}"
    print("  outputs match across placements")

    print(f"\nPart 2: cost of the shape-input copy on a stream with "
          f"{BACKLOG_MATMULS}x{BACKLOG_SIZE}^2 matmuls queued")
    dev_ms, host_ms = copy_cost_under_backlog()
    print(f"  device input  .clone().cpu().to(int64) : {dev_ms:7.2f} ms  (waits for the backlog)")
    print(f"  host input    .contiguous().to(int64)  : {host_ms:7.2f} ms  (does not)")


if __name__ == "__main__":
    main()

Output on stock main:

Part 1: per-call host overhead, idle stream (min of 300 calls)
  shape input on |   host time |     output
--------------------------------------------
            cuda |    157.8 us |   1989.178
             cpu |    196.8 us |   1989.178      <- a host value costs MORE than a device one
  outputs match across placements

Part 2: cost of the shape-input copy on a stream with 20x4096^2 matmuls queued
  device input  .clone().cpu().to(int64) :  132.83 ms  (waits for the backlog)
  host input    .contiguous().to(int64)  :    0.01 ms  (does not)

A host-resident shape value is the slower of the two, because the runtime moves it to the
device and copies it back. Part 2 shows what the copy costs when the stream is busy: it blocks the
host for the full duration of the queued work.

Expected behavior

A shape-tensor input that is already on the host should be used as is, without a device round trip
or a stream synchronization. A device-resident shape input keeps the existing behavior.

Environment

• PyTorch NGC container: 26.07-py3 (torch-tensorrt 2.14.0a0), also reproduces on main
• GPU: NVIDIA L4

[kvnloo]

I like the direction in #4718.

The important performance invariant seems to be:

a host-resident TensorRT shape tensor should remain host-resident unless the engine actually requires a device copy

The current path does:

CPU -> CUDA -> CPU

and the final .cpu() also synchronizes the stream.

For review I would check three cases separately:

• shape tensor already on CPU;
• shape tensor already on CUDA;
• ordinary execution tensor on CUDA.

For the host case, a useful benchmark is not only total inference latency but CPU submission time while unrelated CUDA work is queued on the stream.

If the patch is correct, the host shape path should no longer wait for that queued work merely to recover a value that was already on the host.

The CUDA shape-input path should remain semantically unchanged.