Skip to content

馃悰 [Bug] sort_validator reads an optional argument positionally and raises IndexError#4603

Description

@SrivastavaKshitij

Bug Description

aten::sort(Tensor self, int dim=-1, bool descending=False). When the model writes
torch.sort(x) the traced node carries a single positional argument, but the capability
validator indexes argument 1 unconditionally:

```python
# py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py :: sort_validator
def sort_validator(node: Node, settings: Optional[CompilationSettings] = None) -> bool:
    meta_data = node.args[0].meta.get("tensor_meta")
    if meta_data is None:
        return False
    shape = meta_data.shape
    dim = node.args[1]          # IndexError: tuple index out of range

node.args has length 1, so this raises IndexError. The converter that this validator gates
reads the same argument correctly on the very next screen of the same file:

# aten_ops_sort
return impl.topk.sort(
    ctx, target, SourceIR.ATEN, name,
    args[0],
    dim=args_bounds_check(args, 1, -1),
    descending=args_bounds_check(args, 2, False),
)

The validator is reached through ConverterRegistry.__contains__ -- a membership test. The
partitioner is doing nothing more than asking "could TensorRT take this node?":

# py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py :: is_node_supported
(node in CONVERTERS or node.op == "get_attr")

__contains__ catches only KeyError:

# py/torch_tensorrt/dynamo/conversion/_ConverterRegistry.py :: __contains__
try:
    if isinstance(key, Node):
        self.__getitem__(key)
    else:
        self.__getitem_without_validation__(key)
    return True
except KeyError:
    return False

so the IndexError escapes the membership test and takes the entire compilation with it. The backend's own safety net in
py/torch_tensorrt/dynamo/backend/backends.py is except (AssertionError, RuntimeError, TypeError), which does not include IndexError, so the usual "TRT conversion failed on the
subgraph ... Returning GraphModule forward instead" fallback never runs. We checked: plain
torch.compile(model, backend="tensorrt") with no options at all raises the same IndexError.

To Reproduce

Steps to reproduce the behavior:

Command :

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 traceback

import torch

import torch_tensorrt  # noqa: F401  # registers the "tensorrt" torch.compile backend

COMPILE_OPTIONS = {"pass_through_build_failures": True, "min_block_size": 1}


class SortWithoutDim(torch.nn.Module):
    """Traces to `aten.sort.default(x)` -- one positional argument."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        values, _indices = torch.sort(x)
        return values + 1.0


class SortWithDim(torch.nn.Module):
    """Control: traces to `aten.sort.default(x, 0)` -- two positional arguments."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        values, _indices = torch.sort(x, 0)
        return values + 1.0


def sort_node_arg_count(model: torch.nn.Module, args: tuple[torch.Tensor, ...]) -> int:
    """Returns len(node.args) of the aten.sort.default node in the exported graph."""
    exported = torch.export.export(model, args)
    for node in exported.graph.nodes:
        if node.op == "call_function" and node.target is torch.ops.aten.sort.default:
            return len(node.args)
    raise RuntimeError("no aten.sort.default node in the exported graph")


def compile_and_run(model: torch.nn.Module, args: tuple[torch.Tensor, ...]) -> BaseException | None:
    """Compiles with the TensorRT backend and runs it. Returns the exception, or None.

    Used for the control only. The bug case is run without this wrapper so that its
    exception propagates.
    """
    torch._dynamo.reset()
    optimized = torch.compile(model, backend="tensorrt", options=COMPILE_OPTIONS)
    try:
        optimized(*args)
    except BaseException as exc:  # noqa: BLE001  # the control is expected to succeed
        traceback.print_exc()
        return exc
    return None


def main() -> int:
    print(f"torch {torch.__version__}")
    print(f"torch_tensorrt {torch_tensorrt.__version__}")

    x = torch.rand(4, 8, device="cuda")

    print(
        f"aten.sort.default arg count, torch.sort(x)    : {sort_node_arg_count(SortWithoutDim(), (x,))}"
    )
    print(
        f"aten.sort.default arg count, torch.sort(x, 0) : {sort_node_arg_count(SortWithDim(), (x,))}"
    )

    print("\n=== control: torch.sort(x, 0) ===")
    control_exc = compile_and_run(SortWithDim().cuda().eval(), (x,))
    control_ok = control_exc is None
    print(f"control compiled: {control_ok}")

    print("\n=== bug case: torch.sort(x) ===")
    print("Not wrapped in try/except -- the following call is expected to die with")
    print("    IndexError: tuple index out of range")
    print("raised inside sort_validator in aten_ops_converters.py.\n", flush=True)

    torch._dynamo.reset()
    optimized = torch.compile(
        SortWithoutDim().cuda().eval(), backend="tensorrt", options=COMPILE_OPTIONS
    )
    optimized(x)

    # Only reached on a build where the validator handles the missing argument.
    print("\nbug case compiled cleanly -- the defect is not present on this build")
    print(f"control compiled: {control_ok}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

repro.py

Expected behavior

torch.sort(x) should compile. aten::sort has a default for dim, so the validator should
read the argument the same way the converter does -- args_bounds_check(node.args, 1, -1) --
and then validate normally. torch.sort(x) on a (4, 8) tensor is k = 8, well within the
topk limit the validator is checking for, so it should be accepted and run on TensorRT.

Separately, and independently of this particular validator: a capability validator that raises
should not be able to abort compilation. ConverterRegistry.__contains__ is answering a yes/no
question on behalf of the partitioner; a validator that raises anything other than KeyError
currently escapes it. Treating an exception from a validator as "not supported" (with a warning
naming the validator) would confine any future bug of this shape to a lost fallback opportunity
instead of a failed compile. We would file that as its own request if you prefer.

Environment

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

  • Pytorch NGC container - 26.07-py3

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions