dynamo.compile: nn.Linear on 4D (B,H,W,C) + reshape gives wrong results; identical 3D path is correct

[robbysun]

Bug description

torch_tensorrt.dynamo.compile produces silently wrong numerical results (~23% relative error) when nn.Linear is applied to a 4D tensor (B, H, W, C) and followed by a reshape that splits the channel dimension. The semantically identical computation using a pre-flattened 3D tensor (B, L, C) gives correct results (~0.1% error, just bf16 rounding).

This is a silent correctness bug — no error or warning is raised. We first encountered it while compiling a SAM3 ViT image encoder, where window-attention blocks receive (B*nWindows, ws, ws, C) input: ~50% relative error on features, with detections completely suppressed.

Minimal reproducer

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch_tensorrt

class Attn4D(nn.Module):
    """QKV projection on 4D input — the broken pattern."""
    def __init__(self, C, nH):
        super().__init__()
        self.qkv = nn.Linear(C, C * 3)
        self.proj = nn.Linear(C, C)
        self.nH = nH

    def forward(self, x):                                  # (B, H, W, C)
        B, H, W, C = x.shape; L = H * W; dH = C // self.nH
        qkv = self.qkv(x).reshape(B, L, 3, self.nH, dH)  # ← linear on 4D
        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
        out = F.scaled_dot_product_attention(q, k, v)
        out = out.view(B, self.nH, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, C)
        return self.proj(out)

class Attn3D(nn.Module):
    """Identical math, but input flattened to 3D before linear."""
    def __init__(self, C, nH):
        super().__init__()
        self.qkv = nn.Linear(C, C * 3)
        self.proj = nn.Linear(C, C)
        self.nH = nH

    def forward(self, x):                                  # (B, H, W, C)
        B, H, W, C = x.shape; L = H * W; dH = C // self.nH
        qkv = self.qkv(x.reshape(B, L, C)).reshape(B, L, 3, self.nH, dH)  # ← flatten first
        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
        out = F.scaled_dot_product_attention(q, k, v)
        out = out.view(B, self.nH, L, -1).permute(0, 2, 1, 3).reshape(B, L, C)
        return self.proj(out).reshape(B, H, W, C)

B, H, W, C, nH = 9, 24, 24, 256, 8
device = torch.device("cuda")
torch.manual_seed(0)
x = torch.randn(B, H, W, C, device=device, dtype=torch.bfloat16)

m4 = Attn4D(C, nH).bfloat16().to(device).eval()
m3 = Attn3D(C, nH).bfloat16().to(device).eval()
# Share identical weights so PyTorch outputs are equal
m3.qkv.weight, m3.qkv.bias = m4.qkv.weight, m4.qkv.bias
m3.proj.weight, m3.proj.bias = m4.proj.weight, m4.proj.bias

with torch.no_grad():
    pt4 = m4(x)
    pt3 = m3(x)

print(f"PyTorch 4D vs 3D max-diff: {(pt4 - pt3).abs().max():.2e}")  # ~0

def compile_trt(model, x):
    with torch.no_grad():
        ep = torch.export.export(model, (x,))
    tep = torch_tensorrt.dynamo.compile(
        ep, inputs=[x],
        enabled_precisions={torch.bfloat16},
        use_explicit_typing=False,
        device=torch_tensorrt.Device(gpu_id=0),
    )
    torch_tensorrt.save(tep, "/tmp/repro.ep", inputs=[x])
    return torch_tensorrt.load("/tmp/repro.ep").module()

fn3 = compile_trt(m3, x)
fn4 = compile_trt(m4, x)

with torch.no_grad():
    trt3 = fn3(x)
    trt4 = fn4(x)

ref = pt4.float()
print(f"TRT 3D rel_err: {(trt3.float()-ref).abs().mean()/ref.abs().mean():.5f}")  # ~0.001
print(f"TRT 4D rel_err: {(trt4.float()-ref).abs().mean()/ref.abs().mean():.5f}")  # ~0.23 ← WRONG

Output on RTX 4080:

PyTorch 4D vs 3D max-diff: 0.00e+00
TRT 3D rel_err: 0.00117   ← correct (just bf16 rounding)
TRT 4D rel_err: 0.23456   ← wrong

Additional observations

All individual sub-operations are correct. When nn.Linear, reshape, permute, unbind, RoPE, scaled_dot_product_attention, and the output projection are each compiled and tested in isolation (with pre-computed PyTorch inputs), they all give accurate results (<0.5% error). The wrong result only appears when the full computation is compiled as a single TRT subgraph with a single output.

Returning intermediate tensors as outputs "fixes" it. If the compiled module returns (q, k, v, sdpa_out, final_out) instead of just final_out, all outputs including final_out are accurate. This suggests TRT applies an optimization (possibly a fused MHA kernel) when there is only one output, and that optimized path has a layout bug.

Optimization level makes no difference. Tested optimization_level=0, 1, 2 — all produce the same ~23% error on the 4D path.

Hypothesis. TRT may use NHWC-style internal memory layout for nn.Linear output when the input is a 4D tensor (B, H, W, C). The subsequent reshape(B, L, 3, nH, dH) then misinterprets the memory layout, scrambling the Q/K/V split. The 3D path avoids this because TRT uses standard row-major layout for 3D tensors.

Workaround

Flatten 4D input to 3D before nn.Linear, then reshape back after:

# Before: self.qkv(x)            where x is (B, H, W, C)
# After:  self.qkv(x.reshape(B, H*W, C)).reshape(B, H, W, C_out)

Environment

torch-tensorrt	2.10.0
torch	2.10.0+cu128
tensorrt	10.14.1.48.post1
CUDA	13.0
GPU	NVIDIA GeForce RTX 4080
Python	3.12

Related

• complex_graph_rewrite: AttributeError 'float' object has no attribute 'name' when graph contains scalar literals near view_as_complex #4378 — separate crash in complex_graph_rewrite.py found during the same investigation

Discovered and diagnosed by @robbysun with assistance from Claude Code.

[narendasan]

@robbysun is this still an issue with main. We believe this issue should be resolved in TRT 11

[harshdeepsokhey]

@robbysun is this still an issue ? I am wondering if I can contribute to this ? Thanks!

[robbysun]

@narendasan @harshdeepsokhey Confirmed on TRT 10.14.1.148. I haven't yet tested TRT 11. I'll update the issue once I've had a chance to try TRT 11. Go ahead if you want to contribute. See my workaround above