TensorRT 10.3 FP16 inference produces incorrect segmentation results for DINOv3 ViT models while TensorRT 8.6.1 works correctly #4723

Activity

1. github-actions

   added

   Module:AccuracyOutput mismatch between TensorRT and other frameworksOutput mismatch between TensorRT and other frameworks

   on Mar 16, 2026
2. Whatsonyourmind commented on Apr 4, 2026

   

   Whatsonyourmind

   on Apr 4, 2026

   More actions

   The regression between TRT 8.6.1 and 10.3 for ViT-B16, combined with ViT-L16 failing on all versions, points to FP16 overflow in the attention softmax.

   The numerical chain:

   1. Fused MHA (TRT 8.6.1) vs decomposed attention (TRT 10.3). In TRT 8.6.1's fused MHA kernel, the entire softmax(Q*K^T / sqrt(d_k)) * V computation likely uses FP32 accumulators internally. In TRT 10.3's decomposed path, the MatMul(Q, K^T) output is materialized as an FP16 tensor before being fed to a separate Softmax operator.

   2. Where FP16 breaks. For ViT with d_k=64, scaled attention logits typically range from -10 to +10. The softmax numerator computes exp(z_i - z_max). In FP16:

      • exp(11.09) = 65504 = FP16_MAX
      • Any logit spread > ~11 causes exp() to overflow to Inf

      For ViT-B16 (768-dim, 12 heads), logit spreads occasionally exceed this. For ViT-L16 (1024-dim, 16 heads, 24 layers), deeper layers produce larger activation norms, making overflow nearly guaranteed.

   3. When softmax overflows: numerator = Inf, denominator = Inf, result = NaN. After a few NaN-contaminated attention layers, the output degenerates to a single class.

   4. Why 8.6.1 works for B16 but 10.3 doesn't. The fused kernel computes exp(z_i - z_max) in FP32 internally before writing FP16 output, so the intermediate overflow never happens.

   Recommended fix (minimal precision override):

   config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)
   for i in range(network.num_layers):
       layer = network.get_layer(i)
       if layer.type == trt.LayerType.SOFTMAX:
           layer.precision = trt.float32
           layer.set_output_type(0, trt.float32)

   This forces only Softmax layers to FP32 while keeping everything else in FP16. For ViT-L16, you may also need LayerNorm in FP32, since accumulated normalization errors in deeper networks shift the attention logit distribution.
3. elymsyr commented on Oct 19, 2026

   

   elymsyr

   on Oct 19, 2026

   More actions

   We hit what looks like the same regression with a fine-tuned DINOv2 ViT-B/14 model (keypoint regression head, 392×392 input, 785 tokens) on TensorRT 10.3.0, RTX 3080. The ONNX model was correct in ONNX Runtime, but the TRT engine was badly off: median error ~70 px against ORT.

   It wasn't FP16. The engine was just as wrong in FP32, in strongly-typed builds, and at every builder optimization level we tried. That points away from softmax/LayerNorm overflow, at least in our case. If you can, build your ViT-B16 engine in FP32 and compare with ORT. If FP32 is also wrong, this is probably the same bug.

   Root cause (found by truncating the graph and comparing intermediate outputs with ORT): the PyTorch/timm export splits q/k/v like this:

   qkv [B, N, 3*H*D] -> Reshape [B, N, 3, H, D] -> Split(axis=2) -> Squeeze(axis=2) -> q, k, v [B, N, H, D]
   

   TRT 10.3 (Myelin) computes this Split with the wrong memory layout. In effect it reads [B, N, 3, H, D] as [B, N, H, 3, D], so query head h gets the data of qkv block 3h. Every attention block mixes heads, and the output is garbage, which fits a near-constant segmentation map.

   One debugging trap: marking intermediate tensors as extra network outputs changes Myelin's fusion and makes the bug disappear. Truncate the ONNX graph instead.

   Workaround: rewrite each such pattern in the ONNX before building. Replace Reshape(5D) → Split → Squeeze with three Slice ops on the last axis ([0:HD], [HD:2HD], [2HD:3HD]), each followed by a Reshape to [B, N, H, D]. The math is the same, and ORT output is bit-exact before and after the rewrite. Results with the rewrite:

   Build	Error vs ORT (full-frame px)	Latency
   FP32	max 0.66
   FP16	median 0.03, max 1.55	3.67 ms (3.66 ms unpatched)

   So the FP16 rewrite costs nothing in speed, and FP16 works fine once the layout bug is gone.

   Explicit (non-fused) attention or moving the attention scale did not help us. Only the qkv split rewrite fixed it.

   ViT-L16 also failing on 8.6.1 may be a separate, genuinely FP16-related problem. For the ViT-B16 regression from 8.6.1 to 10.x, though, I'd check for this qkv Split pattern first. I can share our small ONNX rewrite script if it's useful.

   fix_vit_qkv_split.py