🐛 [Bug] engine meta kernel has no ShapeEnv when inputs are statically shaped but outputs are symbolic

[SrivastavaKshitij]

Bug Description

The TensorRT engine meta kernel only ever obtains a ShapeEnv as a side effect of walking
the engine's input shapes. An engine whose inputs are all statically shaped but whose
output shape is symbolic therefore reaches the output loop with shape_env is None and
aborts with an explicit error, even though a perfectly good ShapeEnv is sitting in the
fake_mode the same function already fetched.

py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py :: _apply_symbolic_shape_expressions

    fake_mode = detect_fake_mode(inputs)          # line 37 - fake_mode.shape_env is right here
    ...
    shape_env = None                              # line 55

    # Align inputs: for each captured input, match it with the corresponding runtime input
    for idx, (inp_tensor, inp_info) in enumerate(zip(inputs, input_info)):
        for d, s in zip(inp_tensor.shape, inp_info["shape_exprs"]):
            if isinstance(d, torch.SymInt):
                symbol_to_symint[s] = d
                if shape_env is None:
                    shape_env = d.node.shape_env   # line 63 - the ONLY assignment

and then, in the output loop:

                    elif shape_env is not None:
                        ...
                        output_symint = shape_env.create_symintnode(expr, hint=hint)
                        ...
                    else:
                        raise RuntimeError(                                        # line 149
                            "[torch.ops.tensorrt.execute_engine]: No shape_env available during meta kernel execution"
                        )

shape_env is bound only inside if isinstance(d, torch.SymInt), so "no input dimension is
symbolic" is silently converted into "no shape environment exists". The two things are
unrelated: an engine can mint an unbacked symbol internally (any data-dependent op that has a
converter -- nonzero here -- runs inside the engine and produces a data-dependent output
extent) while every one of its inputs is a fixed size.

Observed error:

RuntimeError: [torch.ops.tensorrt.execute_engine]: No shape_env available during meta kernel execution

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

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

torch._dynamo.config.capture_dynamic_output_shape_ops = True

ROWS = 16


class DataDependentOutput(torch.nn.Module):
    """Statically shaped inputs, data-dependent output row count.

    `nonzero` has a TensorRT converter, so the unbacked symbol for the row count is
    minted at TRT compile time and appears only in the engine's *output* shape
    expressions. Nothing in the engine's input shapes is symbolic.
    """

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        idx = torch.nonzero(mask).squeeze(1)
        return torch.index_select(x, 0, idx) * 2.0


def build_engine() -> torch.fx.GraphModule:
    """Compiles the model to a single TensorRT engine with static input shapes."""
    model = DataDependentOutput().eval().cuda()
    x = torch.randn((ROWS,), device="cuda")
    mask = (torch.arange(ROWS, device="cuda") % 3) == 0
    exported = torch.export.export(model, (x, mask))
    return torch_tensorrt.dynamo.compile(
        exported,
        inputs=(x, mask),
        min_block_size=1,
        pass_through_build_failures=True,
    )


def main(argv: list[str] | tuple[str, ...] = ()) -> int:
    """Drives the compiled engine's meta kernel under a fresh FakeTensorMode."""
    del argv
    print(f"torch          {torch.__version__}")
    print(f"torch_tensorrt {torch_tensorrt.__version__}")

    trt_gm = build_engine()
    print(trt_gm.graph, flush=True)

    # A fresh ShapeEnv + FakeTensorMode is what a re-export (e.g. AOTInductor) sets up.
    # Every fake input is statically shaped, exactly as the engine was compiled.
    shape_env = ShapeEnv()
    reproduced = False
    with FakeTensorMode(shape_env=shape_env):
        x_fake = torch.empty((ROWS,), dtype=torch.float32, device="cuda")
        mask_fake = torch.empty((ROWS,), dtype=torch.bool, device="cuda")
        try:
            out = trt_gm(x_fake, mask_fake)
            print(f"meta kernel returned {out}")
        except Exception as exc:  # pylint: disable=broad-except
            traceback.print_exc()
            reproduced = "No shape_env available during meta kernel execution" in str(exc)

    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]
    %mask : [num_users=1] = placeholder[target=mask]
    %_run_on_acc_0 : [num_users=1] = call_module[target=_run_on_acc_0](args = (%mask, %x), kwargs = {})
    return (_run_on_acc_0,)
Traceback (most recent call last):
  File "/w/repro.py", line 89, in main
    out = trt_gm(x_fake, mask_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>.30", line 6, in forward
    _run_on_acc_0 = self._run_on_acc_0(mask, x);  mask = x = None
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*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 "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py", line 454, in forward
    outputs = torch.ops.tensorrt.execute_engine(input_tensors, self.engine)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 1279, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
[...]
  File "/usr/local/lib/python3.12/dist-packages/torch/library.py", line 1671, in inner
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py", line 219, in fake_tensorrt_execute_engine
    return _apply_symbolic_shape_expressions(inputs, shape_info)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py", line 149, in _apply_symbolic_shape_expressions
    raise RuntimeError(
RuntimeError: [torch.ops.tensorrt.execute_engine]: No shape_env available during meta kernel execution

reproduced: True

Expected behavior

The meta kernel should use the ShapeEnv of the fake mode it is already running under.
fake_mode is bound on line 37, before shape_env is initialised to None, and
fake_mode.shape_env is by construction the environment the surrounding trace uses -- which
is also the environment the newly created output SymInt must belong to for the rest of the
trace to be able to reason about it. Deriving it from an input's d.node.shape_env gets the
same object when an input happens to be symbolic, so seeding from the fake mode is strictly
more general:

    shape_env = fake_mode.shape_env

With that, the raise RuntimeError(...) on line 149 becomes a genuine "should not happen"
guard rather than a routine outcome for any engine with static inputs and a data-dependent
output.

Environment

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

• Pytorch NGC container : 26.07-py3