TRT 11.2: `Transpose` of a network input folded into a fused MHA operand with the permutation applied to the shape but not to the strides

Description

Wrong results produces when a Transpose of a network input is absorbed into a fused MHA operand (_gemm_mha_v2 / _gemv_mha_v1): the fold applies the permutation to the operand's shape but not its strides — the operand is addressed as if the buffer were contiguous in the permuted dimension order, instead of permuting the source tensor's strides.

Details (LLM text):
Boundary (every point rebuilt and re-adjudicated, 245 configurations): fires iff the fused MHA kernel is selected — head dim H ≥ 9 (at H ≤ 8 a non-fused path is taken and the transposes fold correctly, max|diff| exactly 0), batch-of-heads G ≥ 2, any sequence length S from 2 to 16384; float16 and bfloat16 fire, float32 does not (fp32 never selects the fused kernel). Softmax and the second MatMul are both required for fusion. Deterministic: 16/16 independent builds (fresh builder, no shared timing cache) fire with bit-identical wrong outputs. No build-flag workaround found: default settings, DIRECT_IO + allowed_formats=LINEAR on every binding, builder optimization levels 0–5, non-strongly-typed network, and TF32-on all produce the same wrong engine and the same operand strides. The only correct-engine formulation found without changing the input's at-rest layout is listing the transposed tensor as an additional graph output, which forces a real transpose layer (at the cost of materializing the tensor).

Environment

TensorRT Version: 11.2.1.2 (TensorRT v110201); reproduced identically via the python wheel tensorrt-cu12==11.2.1.2 and via the apt packages libnvinfer*=11.2.1.2-1+cuda12.9

NVIDIA GPU: RTX PRO 6000 Blackwell Server Edition

NVIDIA Driver Version: 580.126.20

CUDA Version: runtime 12.9 (nvidia-cuda-runtime-cu12, cudaRuntimeGetVersion() = 12090); driver reports 13000

CUDNN Version: not used (no cuDNN in the reproducer stack)

Operating System: Ubuntu 20.04.6 LTS (x86_64)

Python Version (if applicable): 3.8.10

Tensorflow Version (if applicable): n/a

PyTorch Version (if applicable): n/a (reproducer is pure numpy + onnx; no torch involved)

Baremetal or Container (if so, version): Container (privileged Ubuntu 20.04.6 container, host driver 580.126.20); nothing container-specific — the defect is in build-time operand stride assignment

Relevant Files

Attached repro.py, see reproduce commands

Steps To Reproduce

Commands or scripts:

python3 -m pip install -U pip
python3 -m pip install --extra-index-url https://pypi.nvidia.com tensorrt-cu12==11.2.1.2 nvidia-cuda-runtime-cu12 numpy onnx onnxruntime

python3 repro.py --emit repro.onnx
python3 repro.py repro.onnx --ort-check ; echo "exit=$?"

Expected: exit=0, the engine agrees with the numpy/ONNX-semantics reference to float16 rounding (~1e-3), as onnxruntime does.

Actual: exit=1

1. MAIN -- inputs uploaded exactly as the ONNX declares them
   max|diff| vs reference : 0.185562
   |reference| max        : 0.215576
   ratio                  : 0.8608   (float16 rounding would be ~1e-3)
   elements differing     : 32 / 32

2. CONTROL -- the TRANSPOSED arrays' bytes fed into the same [S,G,H] bindings
   max|diff| vs the SAME reference : 0

Have you tried the latest release?: yes, 11.2.1.2 is the latest release

Can this model run on other frameworks? For example run ONNX model with ONNXRuntime (polygraphy run <model.onnx> --onnxrt): yes, onnxruntime 1.19.2 (CPUExecutionProvider) runs the same repro.onnx and agrees with the numpy reference. repro.py checks this.

Start with the attached repro.py and run its emit and --ort-check commands to confirm the discrepancy against the NumPy/ONNX reference. Trace the fused _gemm_mha_v2 and _gemv_mha_v1 paths, focusing on how a folded Transpose assigns operand shapes and strides. Done means the TensorRT engine matches the reference without requiring the transposed tensor as an additional output.