From f866a0e55896f367925d15795ecc02694d32935c Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Wed, 9 Sep 2026 21:03:49 -0700 Subject: [PATCH 1/2] fix: handle empty one-dimensional tensors in concatenation ## Problem Joining an empty tensor with a larger tensor can force a model to run that operation in PyTorch instead of TensorRT. This happens when a model starts with an empty cache and adds values to it. PyTorch accepts an empty tensor with shape `(0,)` alongside tensors with more dimensions. TensorRT requires matching numbers of dimensions. ## Change Skip these empty tensors when checking dimensions and building the TensorRT concatenation. Continue rejecting dimension mismatches between the remaining tensors. An empty tensor can still affect the result's data type. Choose the output type using every input before removing empty ones. Otherwise, later arithmetic could overflow in float16 when ordinary PyTorch would use float32. The validator's test inputs now include the data type it needs. ## Tests Passed 10/10 focused tests: four validator tests and six tests that build and run TensorRT engines. They cover empty tensors in different positions, positive and negative axes, and mixed input types. They also cover a float64 constant with conversion to float32 enabled. Tested on Linux x86_64 with Python 3.12, TensorRT 11.2, and the native runtime. Windows, aarch64, TensorRT-RTX, TensorRT 10.x, and input dimensions that vary between calls were not tested. --- .../dynamo/conversion/aten_ops_converters.py | 104 +++++++++++++++--- tests/py/dynamo/conversion/test_cat_aten.py | 63 +++++++++++ .../conversion/test_cat_validator_aten.py | 48 ++++++++ 3 files changed, 198 insertions(+), 17 deletions(-) create mode 100644 tests/py/dynamo/conversion/test_cat_validator_aten.py diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 898cd78ed1..80bc656d2e 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -3,6 +3,7 @@ # mypy: disallow-untyped-decorators=False +import functools import logging import operator from typing import ( @@ -22,7 +23,7 @@ import torch from tensorrt import ITensor as TRTTensor from torch.fx.node import Argument, Node, Target -from torch_tensorrt import ENABLED_FEATURES +from torch_tensorrt import ENABLED_FEATURES, _enums from torch_tensorrt._utils import ( is_tensorrt_rtx_version_supported, is_tensorrt_version_supported, @@ -269,16 +270,46 @@ def parse_cat_args( return input_tensors, dim +def _is_rank1_empty_shape(shape: Sequence[int]) -> bool: + """Is this the shape of a rank-1 empty tensor, as torch.tensor([]) produces?""" + return len(shape) == 1 and shape[0] == 0 + + +def _is_rank1_empty(value: Any) -> bool: + """Is this operand a rank-1 empty tensor, as torch.tensor([]) produces?""" + shape = getattr(value, "shape", None) + return shape is not None and _is_rank1_empty_shape(shape) + + +def _promoted_dtype(operands: Sequence[Any]) -> Optional[_enums.dtype]: + """The dtype torch.cat would produce over these operands. + + A rank-1 empty operand contributes nothing to the shape but still takes part in + dtype promotion, so this has to run over the operands before any are dropped. + Returns None when they already agree, leaving the existing behaviour alone. + """ + dtypes = [] + for operand in operands: + dtype = getattr(operand, "dtype", None) + if dtype is None: + return None + dtypes.append(_enums.dtype._from(dtype).to(torch.dtype)) + if not dtypes or all(dtype == dtypes[0] for dtype in dtypes): + return None + return _enums.dtype._from(functools.reduce(torch.promote_types, dtypes)) + + def cat_validator(node: Node, settings: Optional[CompilationSettings] = None) -> bool: """ Validator for torch.cat operation with empty tensor handling. PyTorch allows torch.tensor([]) (shape (0,)) to be concatenated with higher-dimensional - tensors, but TensorRT requires all inputs to have the same rank. This validator catches - this specific edge case. + tensors, but TensorRT requires all inputs to have the same rank. A rank-1 empty operand + holds no elements, so aten_ops_cat leaves it out and the ranks then agree. - Example valid case: cat([(3, 4), (0, 4)], dim=0) - same rank, properly shaped empty tensor for TRT - Example invalid case: cat([(3, 4), (0,)], dim=0) - torch.tensor([]) with rank mismatch + Example valid case: cat([(3, 4), (0,)], dim=0) - the empty operand is dropped + Example invalid case: cat([(0,), (2, 3), (2, 3, 4)], dim=0) - the operands that hold + elements still disagree on rank after the empty one is dropped """ # Use parse_cat_args to properly extract inputs (handles both args and kwargs patterns) inputs, _ = parse_cat_args(node.args, node.kwargs) @@ -288,10 +319,12 @@ def cat_validator(node: Node, settings: Optional[CompilationSettings] = None) -> # Collect metadata for all inputs input_metas = [] + input_dtypes = [] for inp in inputs: if isinstance(inp, TRTTensor): # TRTTensor has shape directly input_metas.append(inp.shape) + input_dtypes.append(inp.dtype) else: # For nodes, get metadata meta = getattr(inp, "meta", {}).get("tensor_meta") @@ -300,6 +333,25 @@ def cat_validator(node: Node, settings: Optional[CompilationSettings] = None) -> return True shape = tuple(meta.shape) input_metas.append(shape) + input_dtypes.append(meta.dtype) + + # Dropping an empty operand also drops it from dtype promotion, so aten_ops_cat works + # the promoted dtype out over every operand and passes it through. TensorRT has no + # float64, so if the promotion lands there and truncate_double is not set, the build + # fails on a graph that used to fall back. Refuse it here instead. + if any(_is_rank1_empty_shape(tuple(shape)) for shape in input_metas): + try: + promoted = functools.reduce(torch.promote_types, input_dtypes) + except TypeError: + promoted = None + if promoted == torch.float64 and not ( + settings is not None and settings.truncate_double + ): + _LOGGER.debug( + "Concatenation rejected by TRT, dropping the empty operand promotes to " + "float64, which needs truncate_double. Falling back to PyTorch" + ) + return False # Check for the specific problematic case: # 1D empty tensor (0,) being concatenated with higher-dimensional tensors @@ -307,18 +359,20 @@ def cat_validator(node: Node, settings: Optional[CompilationSettings] = None) -> # If all ranks are the same, it's fine (PyTorch and TensorRT both handle this) if len(set(ranks)) == 1: return True - # If ranks differ, check if we have a 1D empty tensor (0,) in the mix - # This is the torch.tensor([]) case that PyTorch allows but TensorRT doesn't - for i, shape in enumerate(input_metas): - if shape == (0,) or (len(shape) == 1 and shape[0] == 0): - # Found a 1D empty tensor with rank mismatch - _LOGGER.debug( - f"Concatenation rejected by TRT, torch.tensor([]) or 1D empty tensor at position {i} " - f"PyTorch allows this but TensorRT requires all inputs to have the same rank. " - f"Use torch.empty((0, ...)) with explicit dimensions matching other inputs instead. Falling back to Pytorch" - ) - return False - return True + # A rank-1 empty operand holds no elements, so leaving it out of the + # concatenation gives the same result. Accept the node when the operands that + # do hold elements agree on rank; aten_ops_cat drops the empty ones before + # building the layer. + non_empty = [shape for shape in input_metas if not _is_rank1_empty_shape(shape)] + if len({len(shape) for shape in non_empty}) == 1: + return True + # Reaching here means the operands that survive still disagree on rank, which + # TensorRT cannot concatenate, so name those rather than the empty one. + _LOGGER.debug( + f"Concatenation rejected by TRT, operands {non_empty} do not all have the same " + f"rank. TensorRT requires every operand to have the same rank. Falling back to PyTorch" + ) + return False @dynamo_tensorrt_converter( @@ -334,6 +388,21 @@ def aten_ops_cat( name: str, ) -> Union[TRTTensor, Sequence[TRTTensor]]: inputs, dim = parse_cat_args(args, kwargs) + # TensorRT requires every operand to have the same rank, so a rank-1 empty operand + # has to go. It holds no elements, so leaving it out does not change the result, + # but torch.cat still counts it when it promotes the output dtype: concatenating + # torch.tensor([]), which is float32, onto a float16 tensor produces float32. Work + # the promoted dtype out over every operand and pass it through, or the engine + # builds the concatenation in the survivors' narrower type and overflows. + non_empty = [operand for operand in inputs if not _is_rank1_empty(operand)] + cast_dtype = None + if non_empty and len(non_empty) != len(inputs): + cast_dtype = _promoted_dtype(inputs) + # TensorRT has no float64. The validator only lets float64 reach here when + # truncate_double is set, which means the caller accepts float32, so map it. + if cast_dtype is not None and cast_dtype.to(torch.dtype) == torch.float64: + cast_dtype = _enums.dtype._from(torch.float32) + inputs = non_empty return impl.cat.cat( ctx, target, @@ -341,6 +410,7 @@ def aten_ops_cat( name, input=inputs, dim=dim, + cast_dtype=cast_dtype, ) diff --git a/tests/py/dynamo/conversion/test_cat_aten.py b/tests/py/dynamo/conversion/test_cat_aten.py index b6f5ac0b9a..7aadff5152 100644 --- a/tests/py/dynamo/conversion/test_cat_aten.py +++ b/tests/py/dynamo/conversion/test_cat_aten.py @@ -378,6 +378,69 @@ def forward(self, x): inputs, ) + @parameterized.expand( + [ + ("leading_negative_dim", 0, -2, torch.float32), + ("middle_negative_dim", 1, -2, torch.float32), + ("trailing_positive_dim", 2, 2, torch.float32), + # The empty operand is float32 while the rest are not, so torch.cat + # promotes the result. Dropping the operand must not drop the promotion. + ("promotes_from_float16", 0, -2, torch.float16), + ("promotes_from_int32", 0, -2, torch.int32), + ] + ) + def test_cat_rank1_empty_operand(self, _, position, dim, dtype): + """A rank-1 empty operand should not push the concatenation out of TensorRT. + + transformers emits this on the first write to every layer of a DynamicCache, so + refusing it splits the graph once per layer. + + The arithmetic after the concatenation is deliberate. With the concatenation as + the graph output the engine boundary casts the result back, which hides a wrong + internal dtype; the multiply makes it visible. + """ + + class CatRank1Empty(nn.Module): + def forward(self, x, y): + empty = torch.tensor([], dtype=torch.float32, device=x.device) + operands = [x, y] + operands.insert(position, empty) + return torch.ops.aten.cat.default(tuple(operands), dim) * 1000 + + inputs = [ + torch.full((1, 2, 8, 4), 300, dtype=dtype, device="cuda"), + torch.full((1, 2, 8, 4), 300, dtype=dtype, device="cuda"), + ] + self.run_test( + CatRank1Empty(), + inputs, + use_dynamo_tracer=True, + enable_passes=True, + ) + + def test_cat_rank1_empty_float64_with_truncation(self): + """A float32 tensor concatenated with a float64 empty constant. torch.cat promotes + to float64, and dropping the empty operand carries that promoted dtype through, but + TensorRT has no float64. The converter test harness sets truncate_double, so the + caller accepts float32 and this has to build rather than raise. All operands are + rank 1, so the validator always accepted this shape and TensorRT compiled it before + the empty-operand change.""" + + class CatFloat64Empty(nn.Module): + def forward(self, x): + empty = torch.tensor([], dtype=torch.float64, device=x.device) + # Cast to float32 so the reference is float32 too: with truncate_double the + # engine returns float32, and the test compares dtype as well as values. + return torch.ops.aten.cat.default((empty, x), 0).to(torch.float32) + + inputs = [torch.tensor([0.0, 1.0, 2.0], dtype=torch.float32, device="cuda")] + self.run_test( + CatFloat64Empty(), + inputs, + use_dynamo_tracer=True, + enable_passes=True, + ) + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/conversion/test_cat_validator_aten.py b/tests/py/dynamo/conversion/test_cat_validator_aten.py new file mode 100644 index 0000000000..9436487634 --- /dev/null +++ b/tests/py/dynamo/conversion/test_cat_validator_aten.py @@ -0,0 +1,48 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock + +import torch +from torch.testing._internal.common_utils import run_tests +from torch_tensorrt.dynamo.conversion.aten_ops_converters import cat_validator + + +def _make_cat_node(shapes, dim=0): + """A cat node whose operands carry only the metadata the validator reads.""" + operands = [] + for shape in shapes: + operand = MagicMock() + operand.meta = { + "tensor_meta": SimpleNamespace(shape=torch.Size(shape), dtype=torch.float32) + } + operands.append(operand) + node = MagicMock() + node.args = (operands, dim) + node.kwargs = {} + return node + + +class TestCatValidator(unittest.TestCase): + """Metadata-only checks need no GPU and live outside the converter harness.""" + + def test_rank1_empty_with_agreeing_survivors_is_accepted(self): + """The DynamicCache shape. TensorRT cannot take mixed ranks, but the empty + operand holds nothing, so dropping it leaves operands that agree.""" + self.assertTrue(cat_validator(_make_cat_node([(0,), (1, 2, 8, 4)]))) + self.assertTrue(cat_validator(_make_cat_node([(0,), (2, 3), (5, 3)]))) + + def test_disagreeing_survivors_are_refused(self): + """Dropping the empty operand must not paper over a real mismatch.""" + self.assertFalse(cat_validator(_make_cat_node([(0,), (2, 3), (2, 3, 4)]))) + + def test_uniform_ranks_are_accepted_unchanged(self): + """Nothing is dropped here, so the answer must not depend on this change.""" + self.assertTrue(cat_validator(_make_cat_node([(2, 3), (5, 3)]))) + self.assertTrue(cat_validator(_make_cat_node([(0,), (3,)]))) + + def test_single_operand_is_accepted(self): + self.assertTrue(cat_validator(_make_cat_node([(0,)]))) + + +if __name__ == "__main__": + run_tests() From 6e6518a84f47a9736e3079b61c7a5d38f4d267d3 Mon Sep 17 00:00:00 2001 From: PyTorch Bot Date: Wed, 23 Sep 2026 11:59:29 -0700 Subject: [PATCH 2/2] fix: avoid guards on data-dependent cat lengths The empty-input check can raise when a tensor's length depends on its contents, such as the result of nonzero. This stops compilation before partitioning can finish. Use statically_known_true so an operand is treated as empty only when its length is known to be zero. Add a regression test with data-dependent rank-one operands. Test plan: - New validator test fails before the fix and passes afterward. - 35 concatenation validator and converter tests pass. - Data-dependent compile dry-run passes. - TensorRT execution matches PyTorch for zero, one, three, and five nonzero input elements on Linux x86_64 with PyTorch 2.15 nightly and TensorRT 11.3. --- .../dynamo/conversion/aten_ops_converters.py | 3 ++- .../dynamo/conversion/test_cat_validator_aten.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 80bc656d2e..c37af2b23b 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -22,6 +22,7 @@ import numpy as np import torch from tensorrt import ITensor as TRTTensor +from torch.fx.experimental.symbolic_shapes import statically_known_true from torch.fx.node import Argument, Node, Target from torch_tensorrt import ENABLED_FEATURES, _enums from torch_tensorrt._utils import ( @@ -272,7 +273,7 @@ def parse_cat_args( def _is_rank1_empty_shape(shape: Sequence[int]) -> bool: """Is this the shape of a rank-1 empty tensor, as torch.tensor([]) produces?""" - return len(shape) == 1 and shape[0] == 0 + return len(shape) == 1 and statically_known_true(shape[0] == 0) def _is_rank1_empty(value: Any) -> bool: diff --git a/tests/py/dynamo/conversion/test_cat_validator_aten.py b/tests/py/dynamo/conversion/test_cat_validator_aten.py index 9436487634..8572b1780e 100644 --- a/tests/py/dynamo/conversion/test_cat_validator_aten.py +++ b/tests/py/dynamo/conversion/test_cat_validator_aten.py @@ -43,6 +43,20 @@ def test_uniform_ranks_are_accepted_unchanged(self): def test_single_operand_is_accepted(self): self.assertTrue(cat_validator(_make_cat_node([(0,)]))) + def test_unbacked_rank1_operands_do_not_guard(self): + """Data-dependent lengths must not force a guard during validation.""" + + class M(torch.nn.Module): + def forward(self, x): + a = torch.nonzero(x).flatten().float() + return torch.cat([a, a * 2]) + + ep = torch.export.export(M(), (torch.tensor([0, 1, 1]),), strict=False) + cat = next(n for n in ep.graph.nodes if n.target is torch.ops.aten.cat.default) + for operand in cat.args[0]: + self.assertIsInstance(operand.meta["tensor_meta"].shape[0], torch.SymInt) + self.assertTrue(cat_validator(cat)) + if __name__ == "__main__": run_tests()