Incorrect outputs of TensorRT 10.7 when compiling `F.scaled_dot_product_attention` on GPU L4

[ohadravid]

Description

Compiling an Attention layer which uses torch's scaled_dot_product_attention to a TRT engine results in incorrect outputs.

Environment

Tested against nvcr.io/nvidia/pytorch:24.12-py3 on an L4 GPU

TensorRT Version: 10.7.0.23

NVIDIA GPU: L4

NVIDIA Driver Version:

CUDA Version: 12.6.3

CUDNN Version: 9.6.0.74

Operating System: Ubuntu 24.04

Python Version (if applicable): 3.12

Tensorflow Version (if applicable):

PyTorch Version (if applicable): 2.6.0a0+df5bbc09d1.

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

Relevant Files

Model link:

A gist to reproduce the issue

In high level, we have:

class AttentionUsingScaledDotProduct(nn.Module):
    ...
    def forward(self, x):
        ...
        x = F.scaled_dot_product_attention(
            q,
            k,
            v,
            dropout_p=self.attn_drop.p if self.training else 0.0,
            scale=self.scale,
        )
        ...


class ExplicitAttention(nn.Module):
    ...
    def forward(self, x):
        ...
        q = q * self.scale
        attn = q @ k.transpose(-2, -1)

        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
        ...

And the first version produces an incorrect TRT engine.

This is an image of comparing both ONNX files:
explicit version (left) and the scaled dot product version (right)

Interestingly, patching the bad ONNX to have B=0.125 on the right node (and B=1 on the left) fixes the issue, but not the other way around.

Steps To Reproduce

python min_repro.py

Prints:

Torch: [explicit<->sdpa] Is allclose? True
Torch: [explicit<->mha_fwd] Is allclose? True
Torch: [explicit<->sdpa] Total difference: tensor(0.1525)
Torch: [explicit<->mha_fwd] Total difference: tensor(0.1528)
...
TRT: [explicit<->sdpa] Is allclose? False
TRT: [explicit<->sdpa] Total difference: tensor(1373.3945, device='cuda:0')
TRT: [explicit<->mha_fwd] Is allclose? True
TRT: [explicit<->mha_fwd] Total difference: tensor(0.2232, device='cuda:0')
TRT: Explicit Attention: tensor([-0.0924,  0.0904, -0.3557, -0.4190,  0.0607,  0.2370,  0.0293, -0.2666,
        -0.1009,  0.0399, -0.2203, -0.1261,  0.2094,  0.2609, -0.1429, -0.0457,
        -0.0566, -0.2337, -0.0609,  0.0572,  0.1024, -0.1487,  0.2335, -0.0703,
        -0.2909,  0.0832, -0.1907,  0.0462, -0.3819, -0.2341,  0.1027,  0.2187],
       device='cuda:0')
TRT: Scaled Dot Product Attention: tensor([-0.0920,  0.0905, -0.3552, -0.4191,  0.0605,  0.2372,  0.0291, -0.2665,
        -0.1009,  0.0400, -0.2203, -0.1254,  0.2096,  0.2612, -0.1429, -0.0454,
        -0.0563, -0.2348, -0.0605,  0.0571,  0.1028, -0.1491,  0.2334, -0.0706,
        -0.2912,  0.0824, -0.1899,  0.0453, -0.3816, -0.2335,  0.1030,  0.2190],
       device='cuda:0')
TRT: MHA Forward: tensor([-0.0924,  0.0904, -0.3557, -0.4190,  0.0607,  0.2370,  0.0293, -0.2666,
        -0.1009,  0.0399, -0.2203, -0.1261,  0.2094,  0.2609, -0.1429, -0.0457,
        -0.0566, -0.2337, -0.0609,  0.0572,  0.1024, -0.1487,  0.2335, -0.0703,
        -0.2909,  0.0832, -0.1907,  0.0462, -0.3819, -0.2341,  0.1027,  0.2187],
       device='cuda:0')

Commands or scripts:

Have you tried the latest release?:

Yes

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

Yes, running using onnx runtime works as expected.

[kevinch-nv]

Thanks for the report, able to reproduce on my end. Filing a bug internally to investigate.

[kevinch-nv]

@ohadravid On my end it looks to be resolved in TensorRT 10.8, available in the new nvcr.io/nvidia/pytorch:25.01-py3 container. Can you verify on your end?

[ohadravid]

@kevinch-nv yes! Using the 25.01 container works as expected 🥳

[evolvingai]

Description

I have exported from onnx a Video Transformer model to tensorrt. With tensorrt 10.7 this issue was happening on linux distribution too (cuda 12.2) and when upgrading to tensorrt 10.11 this issue has been fixed.
So, the outputs of the model coming through the transformer part (scaled dot product used) were affected, while the other head (directly from backbone) were ok. AFter upgrading to 10.11 results were correct.

However when moving to windows machine with tensorrt 10.11 the same problem appeared again.

def _multi_head_self_attn(self, x):
"""
x: [B, S, d_model]
returns: [B, S, d_model]
"""
B, S, D = x.shape
qkv = self.qkv_proj(x) # [B, S, 3*D]
qkv = qkv.view(B, S, 3, self.num_heads, self.d_head)
q, k, v = qkv.unbind(dim=2) # each [B, S, H, d_head]
# transpose to [B, H, S, d_head]
q, k, v = [t.permute(0, 2, 1, 3) for t in (q, k, v)]

    # scaled dot-prod: [B, H, S, S]
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_head)
    weights = F.softmax(scores, dim=-1)
    attn   = torch.matmul(weights, v)              # [B, H, S, d_head]
    attn   = attn.permute(0, 2, 1, 3).contiguous()
    attn   = attn.view(B, S, D)                    # [B, S, d_model]
    return self.out_proj(attn)

ONNX Logits [[ 4.624525 -1.0358317 -6.621598 ]], [[9.9651778e-01 3.4691526e-03 1.3012363e-05]], argmax=[0]
Logits output using attention
TRT : [[-0.6519211 0.93750143 -1.7938591 ]], [[0.16076867 0.78791404 0.05131733]], argmax=[1]
Logits output skipping attention
TRT : [[ 3.5375547 -2.3099315 -8.087528 ]], [[9.9711227e-01 2.8788093e-03 8.9132209e-06]], argmax=[0]

Environment

TensorRT Version: 10.11

NVIDIA GPU: A6000

NVIDIA Driver Version: Driver Version: 536.25

CUDA Version: Cuda compilation tools, release 12.2, V12.2.91
Build cuda_12.2.r12.2/compiler.32965470_0

CUDNN Version:

Operating System: Windows

Python Version (if applicable): 3.10

PyTorch Version (if applicable): 2.5.0

Have you tried the latest release?: Latest TensorRT