Incorrect CumSum outputs since TensorRT 10.8+

Description

I've been working on upgrading TensorRT version recently but observed significant accuracy drops for one of our ONNX model. After investigation, we identified the issue is caused by the CumSum layers. Specifically, when multiple CumSum layers coexist inside the model, some CumSum outputs appear incorrect - as if they're operating on the same axis even though they're configured on different axes.

I tried several different TensorRT versions, and the issue affects TensorRT 10.8 and above.

I have already submitted an issue at onnx_tensorrt repo. However, I suspect the root cause may lie within the ICumulativeLayer since constructing a TensorRT network using TensorRT python API still repro the issue (See the test script below).

Environment

TensorRT Version: 10.11.0.33

NVIDIA GPU: A10

NVIDIA Driver Version: 550.144.06

CUDA Version: 12.9

CUDNN Version: 9.10.2

Operating System:

Python Version (if applicable):

Tensorflow Version (if applicable):

PyTorch Version (if applicable):

Baremetal or Container (if so, version): nvcr.io/nvidia/tensorrt::25.06-py3

Relevant Files

Model link:
I can repro for both ONNX models and TRT engines.

• TRT network: Refer to the repro script below.
• ONNX: Check Issue https://github.com/onnx/onnx-tensorrt/issues/1034.

Steps To Reproduce

Commands or scripts:

import tensorrt as trt
import numpy as np
import torch


def create_cumsum_engine(input_shape=(4, 2)):
    """Create TensorRT engine with cumsum operations using add_cumulative layer"""
    logger = trt.Logger(trt.Logger.WARNING)

    builder = trt.Builder(logger)
    network = builder.create_network(0)

    # Input tensor
    input_tensor = network.add_input("data", trt.float32, input_shape)

    # Create scalar axis tensors for cumulative operations (must be 0D/scalar)
    axis_0_tensor = network.add_constant(shape=(), weights=np.array(0, dtype=np.int32))
    axis_1_tensor = network.add_constant(shape=(), weights=np.array(1, dtype=np.int32))

    # Add cumulative layer for axis=0
    cumulative_0 = network.add_cumulative(
        input_tensor,
        axis_0_tensor.get_output(0),
        trt.CumulativeOperation.SUM,
        exclusive=False,
        reverse=False,
    )
    output_0 = cumulative_0.get_output(0)

    # Add cumulative layer for axis=1
    cumulative_1 = network.add_cumulative(
        input_tensor,
        axis_1_tensor.get_output(0),
        trt.CumulativeOperation.SUM,
        exclusive=False,
        reverse=False,
    )
    output_1 = cumulative_1.get_output(0)

    # Mark outputs
    network.mark_output(output_0)
    network.mark_output(output_1)
    output_0.name = "cumsum_axis_0"
    output_1.name = "cumsum_axis_1"

    # Build engine
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 20) # 1 MiB
    serialized_engine = builder.build_serialized_network(network, config)

    runtime = trt.Runtime(logger)
    engine = runtime.deserialize_cuda_engine(serialized_engine)
    return engine


def run_tensorrt_inference(engine, input_tensor: torch.tensor):
    """Run inference using TensorRT engine"""
    context = engine.create_execution_context()

    # Set tensor addresses for v3 API
    context.set_tensor_address("data", input_tensor.data_ptr())

    # Allocate output tensors as torch CUDA tensors
    output_0_tensor = torch.empty_like(input_tensor)
    output_1_tensor = torch.empty_like(input_tensor)

    context.set_tensor_address("cumsum_axis_0", output_0_tensor.data_ptr())
    context.set_tensor_address("cumsum_axis_1", output_1_tensor.data_ptr())

    # Run inference
    context.execute_async_v3(stream_handle=0)

    # Copy outputs back to CPU numpy arrays
    output_0 = output_0_tensor.cpu().numpy()
    output_1 = output_1_tensor.cpu().numpy()

    return output_0, output_1


def test_cumsum_implementations():
    """Test and compare different cumsum implementations"""
    print("Testing cumsum implementations...")

    # Test data
    input_shape = (4, 2)
    test_data = np.ones(input_shape, dtype=np.float32)

    print(f"Input data:\n{test_data}\n")

    # NumPy reference
    np_cumsum_0 = np.cumsum(test_data, axis=0)
    np_cumsum_1 = np.cumsum(test_data, axis=1)

    print("NumPy cumsum axis=0:")
    print(np_cumsum_0)
    print("\nNumPy cumsum axis=1:")
    print(np_cumsum_1)

    # Test TensorRT implementation with add_cumulative
    print("\n" + "=" * 50)
    print("Testing TensorRT add_cumulative approach...")

    engine = create_cumsum_engine(input_shape)
    test_tensor = torch.from_numpy(test_data).float().cuda()
    trt_out_0, trt_out_1 = run_tensorrt_inference(engine, test_tensor)

    print("TensorRT cumsum axis=0:")
    print(trt_out_0)
    print("\nTensorRT cumsum axis=1:")
    print(trt_out_1)


if __name__ == "__main__":
    test_cumsum_implementations()

Output:

> python cumsum_tensorrt_simple.py
Testing cumsum implementations...
Input data:
[[1. 1.]
 [1. 1.]
 [1. 1.]
 [1. 1.]]

NumPy cumsum axis=0:
[[1. 1.]
 [2. 2.]
 [3. 3.]
 [4. 4.]]

NumPy cumsum axis=1:
[[1. 2.]
 [1. 2.]
 [1. 2.]
 [1. 2.]]

==================================================
Testing TensorRT add_cumulative approach...
[08/12/2025-07:03:05] [TRT] [W] Using default stream in enqueueV3() may lead to performance issues due to additional calls to cudaStreamSynchronize() by TensorRT to ensure correct synchronization. Please use non-default stream instead.
TensorRT cumsum axis=0:
[[1. 2.]
 [1. 2.]
 [1. 2.]
 [1. 2.]]

TensorRT cumsum axis=1:
[[1. 2.]
 [1. 2.]
 [1. 2.]
 [1. 2.]]

Have you tried the latest release?: No. The latest version I tried is TRT 10.11.

Can this model run on other frameworks? For example run ONNX model with ONNXRuntime (polygraphy run <model.onnx> --onnxrt):