Repository navigation
Masked bf16 attention rounds its scores to bf16 on Jetson Thor (sm_110), but not on A100 #4938
Description
What is wrong
On Jetson AGX Thor, a bf16 attention layer with a mask (bool or additive float) computes its softmax scores in
bf16. Without a mask it keeps them in float32. On A100 (sm_80), with the same
TensorRT, the masked layer keeps float32 scores, like PyTorch.
When 2 scores are close, bf16 rounding makes them equal, so the softmax weights change. On random inputs the error
grows about 13 times.
What this blocks
On Jetson Thor, evo1 misses its accuracy check and fails (pytorch/executorch#24353).
lingbot_vla_v2 passes there only because LeRobot keeps its attention out of TensorRT on Thor. That splits it into hundreds of pieces and makes it slower than eager PyTorch.
A fix lets both use fused attention on Thor, as they do on the other GPUs.
Repro
The script compiles one scaled_dot_product_attention with Torch-TensorRT, which builds one IAttention layer.
A query and 2 keys are chosen so that their scaled scores (200.45 and 199.55) round to the same bf16 value. With
float32 scores, output channel 0 is +0.42. With bf16 scores it is 0.
import torch, torch.nn.functional as F, torch_tensorrt, tensorrt
B, H, L, S, D = 8, 8, 50, 805, 64
dev = "cuda"
class Attention(torch.nn.Module):
def forward(self, q, k, v, mask=None):
return F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
def near_tie():
q = torch.zeros(B, H, L, D, device=dev); q[..., 0], q[..., 1] = 16.0, 2.0
k = torch.zeros(B, H, S, D, device=dev)
k[..., 0, 0], k[..., 0, 1] = 100.0, 1.796875
k[..., 1, 0], k[..., 1, 1] = 100.0, -1.796875
v = torch.zeros(B, H, S, D, device=dev); v[..., 0, 0], v[..., 1, 0] = 1.0, -1.0
return q.bfloat16(), k.bfloat16(), v.bfloat16()
def random_qkv():
g = torch.Generator(device=dev).manual_seed(0)
return tuple((torch.randn(B, H, n, D, device=dev, generator=g) * s).bfloat16()
for n, s in ((L, 2.0), (S, 3.0), (S, 1.0)))
def reference(q, k, v, mask):
s = q.double() @ k.double().transpose(-2, -1) / D**0.5
if mask is not None:
s = s.masked_fill(~mask, float("-inf")) if mask.dtype == torch.bool else s + mask.double()
return s.softmax(-1) @ v.double()
keep = torch.ones(B, 1, 1, S, dtype=torch.bool, device=dev)
keep[1::2, ..., S - S // 8:] = False # key padding; keys 0 and 1 stay visible
additive = torch.zeros(B, 1, 1, S, device=dev).masked_fill(~keep, float("-inf")).bfloat16().expand(B, H, 1, S).contiguous()
print(torch.cuda.get_device_name(), torch.cuda.get_device_capability(), "TensorRT", tensorrt.__version__)
model = Attention().eval()
for name, mask in (("none", None), ("bool", keep), ("float", additive)):
args = near_tie() + (() if mask is None else (mask,))
trt = torch_tensorrt.dynamo.compile(torch.export.export(model, args), arg_inputs=list(args), min_block_size=1)
with torch.no_grad():
tie = trt(*args)[..., 0].float().mean().item()
rargs = random_qkv() + (() if mask is None else (mask,))
ref = reference(*rargs[:3], mask)
ok = torch.isfinite(ref)
err_trt = (trt(*rargs).double() - ref)[ok].abs().max().item()
err_pt = (model(*rargs).double() - ref)[ok].abs().max().item()
print(f"mask={name:5} channel 0: {tie:+.4f} (float32 scores give +0.4219); "
f"random max error vs float64: TensorRT {err_trt:.4f}, PyTorch {err_pt:.4f}")Results
| no mask | bool mask | float mask | random error, masked | |
|---|---|---|---|---|
| Thor (11,0) | +0.4219 | +0.0000 | +0.0000 | 0.1206 (PyTorch 0.0090) |
| A100 (8,0) | +0.4219 | +0.4219 | +0.4219 | 0.0106 (PyTorch 0.0106) |
Versions: torch 2.15.0 (CUDA 13.2), torch_tensorrt 2.15.0, TensorRT 11.4.0.106 (all from pip install "torch_tensorrt[executorch]"). Thor: JetPack 7.2.1, driver 595.78, CUDA 13.2.
A100: driver 580.126.09.
Actual
Jetson AGX Thor, driver 595.78:
NVIDIA Thor (11, 0) TensorRT 11.4.0.106
mask=none channel 0: +0.4219 (float32 scores give +0.4219); random max error vs float64: TensorRT 0.0090, PyTorch 0.0084
mask=bool channel 0: +0.0000 (float32 scores give +0.4219); random max error vs float64: TensorRT 0.1206, PyTorch 0.0090
mask=float channel 0: +0.0000 (float32 scores give +0.4219); random max error vs float64: TensorRT 0.1206, PyTorch 0.0090
A100 80 GB, driver 580.126.09:
NVIDIA PG509-210 (8, 0) TensorRT 11.4.0.106
mask=none channel 0: +0.4219 (float32 scores give +0.4219); random max error vs float64: TensorRT 0.0106, PyTorch 0.0106
mask=bool channel 0: +0.4219 (float32 scores give +0.4219); random max error vs float64: TensorRT 0.0106, PyTorch 0.0106
mask=float channel 0: +0.4219 (float32 scores give +0.4219); random max error vs float64: TensorRT 0.0106, PyTorch 0.0106
Expected
Float32 scores for masked attention on Thor, as on A100.
Activity