Repository navigation
IAttention gives wrong fp16 and bf16 results when the key is a constant #4935
Description
What is wrong
When the key of an IAttention layer is a constant (for example a frozen weight or a cached key), TensorRT 11.4
computes fp16 and bf16 attention wrongly. The largest error against a float32 reference is about 0.81, where
about 0.001 is expected. There is no error or warning.
The same network is correct when the key is a network input, and it is correct in fp32. A transpose in front of the
key makes no difference. We found it through Torch-TensorRT, which now keeps such attention out of TensorRT, and it
reproduces with the TensorRT Python API alone. It may be related to #4826 (a transpose in front of an attention
operand), but here the transpose is not needed.
What this blocks
No policy in pytorch/executorch#24353 gets wrong results today, because Torch-TensorRT 2.15 does not use TensorRT's fused attention when the key is a constant (pytorch/TensorRT#4964).
Models with such attention run it outside TensorRT. Anyone who uses the TensorRT API directly gets wrong fp16 and bf16 results with no error.
A fix gives these models correct fused attention.
Repro
Needs tensorrt and torch (torch only for buffers and the reference). --key const|input picks a constant or a
network input for the key, --no-transpose stores the key in its final layout.
import argparse
import numpy as np
import tensorrt as trt
import torch
p = argparse.ArgumentParser()
p.add_argument("--dtype", default="fp16", choices=["fp16", "bf16", "fp32"])
p.add_argument("--key", default="const", choices=["const", "input"])
p.add_argument("--no-transpose", action="store_true")
a = p.parse_args()
H, L, S, D = 8, 64, 128, 64
tdt, pdt = {"fp16": (trt.float16, torch.float16), "bf16": (trt.bfloat16, torch.bfloat16),
"fp32": (trt.float32, torch.float32)}[a.dtype]
g = torch.Generator().manual_seed(0)
q = (torch.randn(1, H, L, D, generator=g) * D**-0.5).to(pdt).cuda() # IAttention applies no 1/sqrt(D) scale itself
v = torch.randn(1, H, S, D, generator=g).to(pdt).cuda()
k_stored = torch.randn(1, S, H, D, generator=g).to(pdt) # stored as [1, S, H, D], read as [1, H, S, D]
k = k_stored.transpose(1, 2).contiguous()
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
net = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
tq = net.add_input("q", tdt, tuple(q.shape))
tv = net.add_input("v", tdt, tuple(v.shape))
src = k if a.no_transpose else k_stored
if a.key == "const":
raw = src.view(torch.int16).numpy() if a.dtype == "bf16" else src.numpy()
w = trt.Weights(tdt, raw.ctypes.data, raw.size)
tk = net.add_constant(tuple(src.shape), w).get_output(0)
else:
tk = net.add_input("k", tdt, tuple(src.shape))
if not a.no_transpose:
shuffle = net.add_shuffle(tk)
shuffle.first_transpose = trt.Permutation([0, 2, 1, 3])
tk = shuffle.get_output(0)
att = net.add_attention_v2(tq, tk, tv, trt.AttentionNormalizationOp.SOFTMAX, trt.CausalMaskKind.NONE)
att.decomposable = True
net.mark_output(att.get_output(0))
engine = trt.Runtime(logger).deserialize_cuda_engine(
builder.build_serialized_network(net, builder.create_builder_config()))
ctx = engine.create_execution_context()
out = torch.empty(1, H, L, D, dtype=pdt, device="cuda")
bufs = {"q": q, "v": v, "k": src.cuda().contiguous()}
for i in range(engine.num_io_tensors):
name = engine.get_tensor_name(i)
ctx.set_tensor_address(name, (out if name not in bufs else bufs[name]).data_ptr())
ctx.execute_async_v3(torch.cuda.current_stream().cuda_stream)
torch.cuda.synchronize()
ref = torch.softmax(q.float() @ k.float().cuda().transpose(-2, -1), -1) @ v.float()
print(trt.__version__, a, "max error", (out.float() - ref).abs().max().item())Results
TensorRT 11.4.0.106, max error against the float32 reference. H100 and A100 agree except fp16: constant key 0.8069 on
H100 and 0.8075 on A100, input key 0.0009 and 0.0013.
| dtype | key | with transpose | without transpose |
|---|---|---|---|
| fp16 | constant | 0.8069 | 0.8069 |
| fp16 | input | 0.0009 | 0.0009 |
| bf16 | constant | 0.8084 | 0.8084 |
| bf16 | input | 0.0019 | 0.0019 |
| fp32 | constant | 0.0004 | 0.0004 |
| fp32 | input | 0.0004 | 0.0004 |
Actual
H100, driver 580.126.09:
11.4.0.106 Namespace(dtype='fp16', key='const', no_transpose=False) max error 0.8068886995315552
11.4.0.106 Namespace(dtype='fp16', key='const', no_transpose=True) max error 0.8068886995315552
11.4.0.106 Namespace(dtype='fp16', key='input', no_transpose=False) max error 0.0008955597877502441
11.4.0.106 Namespace(dtype='bf16', key='const', no_transpose=False) max error 0.8084346652030945
11.4.0.106 Namespace(dtype='bf16', key='input', no_transpose=False) max error 0.001896500587463379
11.4.0.106 Namespace(dtype='fp32', key='const', no_transpose=False) max error 0.00039896368980407715
Expected
The same result for a constant key as for an input key, about 0.001 in fp16 and bf16.
Versions
TensorRT 11.4.0.106 (pip), torch 2.15.0 (CUDA 13.2, buffers and reference only), Linux x86_64. H100 and A100 80 GB,
driver 580.126.09.
Activity