Cast/normalize attn_bias in scaled_dot_product_efficient_attention

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

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()