From dd86338d2700325208d67b029462e2bb29255df4 Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Tue, 18 Aug 2026 20:11:18 +0000 Subject: [PATCH 1/2] Fix cumsum dtype --- .../dynamo/conversion/aten_ops_converters.py | 1 + .../dynamo/conversion/impl/slice/ops.py | 26 ++++++++++++++++--- .../py/dynamo/conversion/test_cumsum_aten.py | 23 ++++++++++++++++ 3 files changed, 47 insertions(+), 3 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 56ad2512e0..0b4848c537 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -1427,6 +1427,7 @@ def aten_ops_cumsum( name, args[0], args[1], + kwargs.get("dtype"), ) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py index 90153e24b6..e43ad41965 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py @@ -4,13 +4,16 @@ import numpy as np import tensorrt as trt +import torch from tensorrt import ITensor as TRTTensor from torch.fx.node import Target +from torch_tensorrt import _enums from torch_tensorrt.dynamo._SourceIR import SourceIR from torch_tensorrt.dynamo.conversion import impl from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext from torch_tensorrt.dynamo.conversion.converter_utils import ( calculate_strides, + cast_trt_tensor, flatten_dims, get_positive_dim, get_trt_tensor, @@ -366,7 +369,21 @@ def cumsum( name: str, input: TRTTensor, dim: int, + dtype: Optional[torch.dtype] = None, ) -> TRTTensor: + # aten widens integer cumsum to int64 unless dtype is set; floats keep their type. + input_dtype = _enums.dtype._from(input.dtype).to(torch.dtype) + if dtype is not None: + acc_dtype = dtype + elif not input_dtype.is_floating_point: + acc_dtype = torch.int64 + else: + acc_dtype = input_dtype + acc_np_dtype = _enums.dtype._from(acc_dtype).to(np.dtype) + + if input_dtype != acc_dtype: + input = cast_trt_tensor(ctx, input, acc_dtype, f"{name}_input_cast") + input_shape = input.shape dim = get_positive_dim(dim, len(input_shape)) if input_shape[dim] < 0: @@ -400,11 +417,11 @@ def cumsum( else: data_shape.append(input_shape[i]) zero_trttensor = impl.full.full( - ctx, target, source_ir, name + "_full", data_shape, 0.0 + ctx, target, source_ir, name + "_full", data_shape, 0, dtype=acc_dtype ) else: new_dims = tuple(data.shape) - zeros = np.zeros(new_dims, dtype=np.float32) + zeros = np.zeros(new_dims, dtype=acc_np_dtype) zero_trttensor = get_trt_tensor(ctx, zeros, f"{name}_initial_value") running_sum = loop.add_recurrence(zero_trttensor) @@ -424,7 +441,10 @@ def cumsum( loop_output = loop.add_loop_output(current_sum, trt.LoopOutput.CONCATENATE, dim) set_layer_name(loop_output, target, f"{name}_loop_output", source_ir) loop_output.set_input(1, trip_limit) - return loop_output.get_output(0) + out = loop_output.get_output(0) + if _enums.dtype._from(out.dtype).to(torch.dtype) != acc_dtype: + out = cast_trt_tensor(ctx, out, acc_dtype, f"{name}_output_cast") + return out def tile( diff --git a/tests/py/dynamo/conversion/test_cumsum_aten.py b/tests/py/dynamo/conversion/test_cumsum_aten.py index 1a83c9c45f..919e148095 100644 --- a/tests/py/dynamo/conversion/test_cumsum_aten.py +++ b/tests/py/dynamo/conversion/test_cumsum_aten.py @@ -98,6 +98,29 @@ def forward(self, x): immutable_weights=False, ) + @parameterized.expand( + [ + (torch.int32, torch.int32), + (torch.int32, None), + (torch.float16, None), + (torch.float32, torch.float16), + ] + ) + def test_cumsum_dtype(self, input_dtype, out_dtype): + class Cumsum(nn.Module): + def forward(self, x): + if out_dtype is None: + return torch.ops.aten.cumsum.default(x, 0) + return torch.ops.aten.cumsum.default(x, 0, dtype=out_dtype) + + inputs = [torch.tensor([1, 2, 3, 4], dtype=input_dtype)] + self.run_test( + Cumsum(), + inputs, + use_dynamo_tracer=True, + immutable_weights=False, + ) + if __name__ == "__main__": run_tests() From b398035852671788b4f6b0f503b59c078703be92 Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Fri, 21 Aug 2026 22:36:38 +0000 Subject: [PATCH 2/2] More test coverage --- .../dynamo/conversion/impl/slice/ops.py | 21 +++---- .../py/dynamo/conversion/test_cumsum_aten.py | 55 ++++++++++++++++++- 2 files changed, 62 insertions(+), 14 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py index e43ad41965..2a41367491 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py @@ -371,7 +371,8 @@ def cumsum( dim: int, dtype: Optional[torch.dtype] = None, ) -> TRTTensor: - # aten widens integer cumsum to int64 unless dtype is set; floats keep their type. + # aten.cumsum accumulates bool and integer inputs in int64 and floats in + # their own dtype; an explicit dtype wins over both input_dtype = _enums.dtype._from(input.dtype).to(torch.dtype) if dtype is not None: acc_dtype = dtype @@ -379,7 +380,6 @@ def cumsum( acc_dtype = torch.int64 else: acc_dtype = input_dtype - acc_np_dtype = _enums.dtype._from(acc_dtype).to(np.dtype) if input_dtype != acc_dtype: input = cast_trt_tensor(ctx, input, acc_dtype, f"{name}_input_cast") @@ -416,13 +416,13 @@ def cumsum( ) else: data_shape.append(input_shape[i]) - zero_trttensor = impl.full.full( - ctx, target, source_ir, name + "_full", data_shape, 0, dtype=acc_dtype - ) else: - new_dims = tuple(data.shape) - zeros = np.zeros(new_dims, dtype=acc_np_dtype) - zero_trttensor = get_trt_tensor(ctx, zeros, f"{name}_initial_value") + data_shape = list(data.shape) + + # full rather than np.zeros: numpy has no bf16 + zero_trttensor = impl.full.full( + ctx, target, source_ir, f"{name}_initial_value", data_shape, 0, dtype=acc_dtype + ) running_sum = loop.add_recurrence(zero_trttensor) set_layer_name(running_sum, target, f"{name}_running_sum", source_ir) @@ -441,10 +441,7 @@ def cumsum( loop_output = loop.add_loop_output(current_sum, trt.LoopOutput.CONCATENATE, dim) set_layer_name(loop_output, target, f"{name}_loop_output", source_ir) loop_output.set_input(1, trip_limit) - out = loop_output.get_output(0) - if _enums.dtype._from(out.dtype).to(torch.dtype) != acc_dtype: - out = cast_trt_tensor(ctx, out, acc_dtype, f"{name}_output_cast") - return out + return loop_output.get_output(0) def tile( diff --git a/tests/py/dynamo/conversion/test_cumsum_aten.py b/tests/py/dynamo/conversion/test_cumsum_aten.py index 919e148095..4cee79febd 100644 --- a/tests/py/dynamo/conversion/test_cumsum_aten.py +++ b/tests/py/dynamo/conversion/test_cumsum_aten.py @@ -100,10 +100,13 @@ def forward(self, x): @parameterized.expand( [ - (torch.int32, torch.int32), - (torch.int32, None), + (torch.int32, torch.int32), # explicit dtype keeps int32 + (torch.int32, None), # integral promotes to int64 + (torch.int64, None), (torch.float16, None), + (torch.bfloat16, None), (torch.float32, torch.float16), + (torch.float32, torch.bfloat16), ] ) def test_cumsum_dtype(self, input_dtype, out_dtype): @@ -113,6 +116,7 @@ def forward(self, x): return torch.ops.aten.cumsum.default(x, 0) return torch.ops.aten.cumsum.default(x, 0, dtype=out_dtype) + # 1,2,3,4 accumulate exactly in every dtype under test inputs = [torch.tensor([1, 2, 3, 4], dtype=input_dtype)] self.run_test( Cumsum(), @@ -121,6 +125,53 @@ def forward(self, x): immutable_weights=False, ) + @parameterized.expand( + [ + (torch.int32, None), + (torch.int64, None), + (torch.float32, torch.int64), + ] + ) + def test_cumsum_accumulator_is_exact(self, input_dtype, out_dtype): + class Cumsum(nn.Module): + def forward(self, x): + if out_dtype is None: + return torch.ops.aten.cumsum.default(x, 0) + return torch.ops.aten.cumsum.default(x, 0, dtype=out_dtype) + + # 2**24+1 is unrepresentable in float32, so a float accumulator stalls + # here while an integral one keeps counting; the sums must be exact + inputs = [torch.tensor([2**24, 1, 1, 1], dtype=input_dtype)] + self.run_test( + Cumsum(), + inputs, + rtol=0, + atol=0, + use_dynamo_tracer=True, + immutable_weights=False, + ) + + @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) + def test_cumsum_dynamic_shape_dtype(self, input_dtype): + class Cumsum(nn.Module): + def forward(self, x): + return torch.ops.aten.cumsum.default(x, 0) + + # a dynamic non-cumsum dim sends the seed down full's shape-tensor path + inputs = [ + torch_tensorrt.Input( + min_shape=(1, 2), + opt_shape=(2, 3), + max_shape=(3, 4), + dtype=input_dtype, + ), + ] + self.run_test_with_dynamic_shape( + Cumsum(), + inputs, + immutable_weights=False, + ) + if __name__ == "__main__": run_tests()