Skip to content

Silently incorrect fp16/bf16 outputs from fused MHA on consumer Blackwell (sm_120) for ViT/DETR model (SAM 3 detector) — fp32 exact, persists in 10.16.1, breaking the MHA fusion restores accuracy #4837

Description

@crisreimberg

Description

fp16 and bf16 engines built from a ViT-based DETR-style detector (Meta SAM 3 image detector: ViT backbone at 1008×1008 → transformer encoder → decoder, text prompt "person" baked as constants, outputs boxes cxcywh + scores, no masks) produce silently wrong outputs on a consumer Blackwell GPU (sm_120). There is no error, no warning; trtexec prints PASSED. On a frame with 3 people, the top score drops from 0.933 to 0.045 and 0 detections survive a 0.5 threshold; on an empty frame the 10.16 engine hallucinates 4 detections. The fp32 engine is exact (matches PyTorch: 15/15 boxes over 8 frames, max score delta 0.018, max box delta 0.3 px).

PyTorch itself is correct in fp16 and bf16 (autocast matches fp32 within 0.005/0.021 max score delta), so this is not an inherent precision limitation of the model.

Bisection: the fused multi-head attention is the culprit

Using polygraphy with intermediate tensors marked (reference = the exact fp32 engine):

  1. Marking ~27 boundary tensors: patch_embed output PASSES (abs err 0.012), the first ViT block output already FAILS (abs err 7.2), error grows monotonically through the 32 blocks (up to ~330).
  2. Marking 16 tensors inside block 0: everything PASSES (0.006–0.09) — marking outputs prevents fusion, so the defect lives in a fusion.
  3. Breaking one fusion at a time: breaking only the attention (marking the attn@v output) drops the block-0 error from 7.22 to 0.089. Breaking MLP/GELU or LayerNorm changes nothing.

The block-0 attention shape is: 9 windows × 16 heads × 576 tokens × head dim 64 (windowed ViTDet attention; 4 global blocks use 5184 tokens).

What does NOT fix it

  • --bf16 instead of --fp16 — same failure.
  • Strongly-typed engines built from ONNX exported under torch.autocast (explicit casts recorded; PyTorch evaluates the same graph correctly) — same failure.
  • --precisionConstraints=obey --layerPrecisions=*Softmax:fp32 — the log confirms Set layer .../attn/Softmax to precision fp32 for all 68 softmaxes, yet the result is unchanged (and still fast), i.e. the MHA fusion appears to override / ignore the per-layer precision constraint.
  • Upgrading: wrong on TRT 10.14.1.48 (DeepStream 9.0), 10.16.1.11 (DeepStream 9.1) and 10.16.1.48 (nvcr.io/nvidia/tensorrt:26.04-py3).

What DOES fix it (workarounds, both ~1.8× slower than the fused engine)

  • Rewriting softmax manually in the exported graph (exp(s − max) / sum, no Softmax op) → 15/15 boxes correct in fp16.
  • Keeping Softmax but inserting a numerically neutral Clip(p, 0, 1) between the softmax and the second MatMul (breaks the MatMul → Softmax → MatMul pattern) → 15/15 boxes correct in fp16.

Both confirm the unfused fp16 math is fine on this GPU; only the fused MHA kernel path is wrong.

Measurements (RTX 5070 Ti Laptop, batch 1, 1008×1008, --useCudaGraph --noDataTransfers)

Build Latency (min · median) Boxes vs PyTorch fp32 (15 boxes / 8 frames)
fp32 523 · 546 ms 15/15 correct
fp16 (TRT 10.14) 116 · 128 ms 0/15 wrong
bf16 (TRT 10.14) 163 · 179 ms 0/15 wrong
fp16 strongly-typed (autocast export) 118 · 129 ms 0/15 wrong
fp16 + *Softmax:fp32 constraint 118 · 131 ms 0/15 wrong
fp16 (TRT 10.16.1.11 / DeepStream 9.1) 113 · 122 ms 0/15 wrong + hallucinations
fp16 (TRT 10.16.1.48 / tensorrt:26.04-py3) 106 · 113 ms 0/15 wrong + hallucinations
fp16 with MHA fusion broken (Clip no-op) 211 · 227 ms 15/15 correct

Likely related report

TensorRT 10.14 silently produces wrong detection scores for D-FINE (DETR-style) models on RTX 5090 / sm_120 — same symptom family on sm_120. That reporter says 10.16 / DeepStream 9.1 fixed their model; it does not fix this one (verified with the same DS 9.1 container).

Environment

TensorRT Version: 10.14.1.48, 10.16.1.11, 10.16.1.48 (all reproduce)

NVIDIA GPU: GeForce RTX 5070 Ti Laptop (Blackwell consumer, sm_120, CC 12.0)

NVIDIA Driver Version: 596.21

CUDA Version: 13.2 (containers: 13.0/13.2)

CUDNN Version: as shipped in the containers below

Operating System: Windows 11 + WSL2 2.7.3 + Docker Desktop (containers: nvcr.io/nvidia/deepstream:9.0-samples-multiarch, deepstream:9.1-samples-multiarch, tensorrt:26.04-py3)

Python Version (if applicable): 3.12

PyTorch Version (if applicable): 2.10.0+cu128 (export only; ONNX opset 17, TorchScript exporter)

Baremetal or Container (if so, version): Containers above

Steps To Reproduce

Attached repro_sam3_mha.zip contains: one preprocessed input (input_1x3x1008x1008_fp32.bin), the PyTorch fp32 reference outputs, the correct fp32-engine outputs, the wrong fp16-engine outputs from 10.14 and 10.16.1.11, a compare.py, and the export script.

trtexec --onnx=sam3_person_b1.onnx --saveEngine=fp16.plan --fp16
trtexec --loadEngine=fp16.plan --loadInputs=image:input_1x3x1008x1008_fp32.bin \
        --exportOutput=out.json --warmUp=0 --iterations=1 --duration=0
python compare.py out.json     # -> 0 matched, max |score delta| ~0.93

The ONNX is 1.9 GB (fp32, single input image 1×3×1008×1008, outputs boxes [1,200,4] cxcywh-normalized and scores [1,200], opset 17). I can share it via a download link on request, or it can be regenerated from the public gated checkpoint facebook/sam3 (Hugging Face) with the attached export_sam3_detector_onnx.py (detector-only export: real-valued RoPE, plain MLP instead of the repo's fused _addmm_activation, zero-size prompt tensors replaced — all validated to match the original model box-for-box in PyTorch).

Have you tried the latest release?: Yes — wrong on 10.16.1.48 (26.04 container).

Can this model run on other frameworks?: Yes — ONNX Runtime (fp32) and PyTorch (fp32/fp16/bf16) all produce correct results; the TensorRT fp32 engine as well.

repro_sam3_mha.zip

Metadata

Metadata

Assignees

No one assigned

    Labels

    Module:AccuracyOutput mismatch between TensorRT and other frameworks

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions