diff --git a/py/torch_tensorrt/dynamo/conversion/impl/attention.py b/py/torch_tensorrt/dynamo/conversion/impl/attention.py index 277df03db5..15e711d848 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/attention.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/attention.py @@ -430,6 +430,17 @@ def scaled_dot_product_efficient_attention( # TRT's IAttention layer does not support passing in both attn_bias/mask and causal mask at the same time, # so we convert causal mask to an additive causal mask and add it to the attn_bias attn_bias = get_trt_tensor(ctx, attn_bias, f"{name}_attn_bias") + # attn_bias is combined with an additive (query.dtype) causal mask below via + # elementwise add, so it must already be a matching-dtype additive bias. + if attn_bias.dtype != query.dtype: + attn_bias = cast_trt_tensor( + ctx, + attn_bias, + query.dtype, + f"{name}_cast_attn_bias", + target, + source_ir, + ) L = impl.shape.shape(ctx, target, source_ir, f"{name}_L", query, -2) S = impl.shape.shape(ctx, target, source_ir, f"{name}_S", key, -2) @@ -469,7 +480,23 @@ def scaled_dot_product_efficient_attention( else: if attn_bias is not None: attn_bias = get_trt_tensor(ctx, attn_bias, f"{name}_attn_bias") - attention_layer.mask = attn_bias + if attn_bias.dtype == trt.DataType.BOOL: + mask = attn_bias + elif attn_bias.dtype != query.dtype: + mask = cast_trt_tensor( + ctx, + attn_bias, + query.dtype, + f"{name}_cast_attn_bias", + target, + source_ir, + ) + else: + mask = attn_bias + mask = _normalize_attention_mask_rank( + ctx, mask, query, f"{name}_normalize_attn_bias" + ) + attention_layer.mask = mask fp8_norm = _maybe_set_fp8_softmax(ctx, name, attention_layer) attention_layer.decomposable = not fp8_norm diff --git a/tests/py/dynamo/conversion/test_attention.py b/tests/py/dynamo/conversion/test_attention.py index 387f44065f..2117d370c9 100644 --- a/tests/py/dynamo/conversion/test_attention.py +++ b/tests/py/dynamo/conversion/test_attention.py @@ -101,6 +101,45 @@ def forward(self, query, key, value): decompose_attention=True, ) + @parameterized.expand([((1, 2, 16, 32), (1, 2, 16, 32))]) + def test_sdpa_no_causal_with_bias(self, query_shape, key_shape): + class SDPA(nn.Module): + def forward(self, query, key, value, attn_bias): + attn = torch.ops.aten._scaled_dot_product_efficient_attention.default( + query, + key, + value, + attn_bias, + False, + 0, + False, # is_causal + scale=0.5, + ) + return attn[0] + + inputs = [] + query = torch.randn(query_shape, dtype=torch.float16) + key = torch.rand(key_shape, dtype=torch.float16) + value = torch.rand(key_shape, dtype=torch.float16) + # Regression test: an attn_bias whose dtype doesn't match query's dtype + # (e.g. an int32 padding mask, as HF BERT passes) used to be handed + # straight to TensorRT's IAttention layer with no cast, which TRT's + # AttentionInput layer rejects for anything but Float/Half/BFloat16/ + # Bool matching the attention dtype. int32 (not e.g. float32) is used + # here since PyTorch's own eager kernel accepts an int32 bias but + # rejects a mismatched float dtype outright. + attn_bias = torch.zeros((query_shape[0], 1, 1, key_shape[2]), dtype=torch.int32) + inputs.extend([query, key, value, attn_bias]) + self.run_test( + SDPA(), + inputs, + rtol=1e-2, + atol=1e-2, + precision=torch.float16, + enable_passes=True, + decompose_attention=True, + ) + if __name__ == "__main__": run_tests()