fix: optimize attention decomposition

py/torch_tensorrt/dynamo/conversion/impl/arange.py

@@ -24,12 +24,12 @@ def arange(
     step: Union[int, TRTTensor],
 ) -> TRTTensor:
     """
-    Creates a sequence of values (arange) either dynamically or statically,
-    then outputs a TensorRT tensor.
+    Creates a sequence of values (arange) with a TensorRT Fill layer.
 
-    If any of (start, end, step) is a TRT tensor, it sets up a dynamic arange
-    using a Fill layer. Otherwise, it creates a static NumPy array and converts
-    it into a TensorRT constant tensor.
+    If any of (start, end, step) is a TRT tensor, the Fill output length is
+    computed dynamically. Otherwise, NumPy is used only to determine the static
+    output length. Keeping static ranges as Fill layers preserves their sequence
+    provenance for downstream TensorRT graph-pattern recognition.
     """
     # If any argument is a TRT tensor, use dynamic arange with a Fill layer
     if any(isinstance(x, TRTTensor) for x in (start, end, step)):
@@ -68,9 +68,20 @@ def arange(
         return fill_layer.get_output(0)
 
     else:
-        # All arguments are static, so use NumPy arange and create a TRT constant
-        arr = np.arange(start, end, step, dtype=np.int32)
-        weights = trt.Weights(arr)
-        const_layer = ctx.net.add_constant(arr.shape, weights)
-        set_layer_name(const_layer, target, f"{name}_arange_const", source_ir)
-        return const_layer.get_output(0)
+        # Keep a static arange as LINSPACE rather than materializing its values in
+        # a Constant. TensorRT uses this producer provenance when recognizing
+        # compact causal attention masks.
+        output_shape = np.arange(start, end, step, dtype=np.int32).shape
+        start_tensor = get_trt_tensor(
+            ctx, start, name + "_start", dtype=trt.int32, min_rank=0
+        )
+        step_tensor = get_trt_tensor(
+            ctx, step, name + "_step", dtype=trt.int32, min_rank=1
+        )
+        fill_layer = ctx.net.add_fill(
+            output_shape, trt.FillOperation.LINSPACE, trt.int32
+        )
+        fill_layer.set_input(1, start_tensor)
+        fill_layer.set_input(2, step_tensor)
+        set_layer_name(fill_layer, target, f"{name}_arange_fill", source_ir)
+        return fill_layer.get_output(0)

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -485,18 +485,21 @@ def scaled_dot_product_attention_decomposition(
     L, S = query.size(-2), key.size(-2)
     device = query.device
 
-    if is_causal or attn_mask is not None:
-        attn_bias = torch.zeros((L, S), dtype=query.dtype, device=device)
-
     if is_causal:
         assert attn_mask is None, "attn_mask must be None when is_causal=True"
+        attn_bias = torch.zeros((L, S), dtype=query.dtype, device=device)
         temp_mask = torch.ones((L, S), dtype=torch.bool, device=device).tril(diagonal=0)
         attn_bias = attn_bias.masked_fill(temp_mask.logical_not(), float("-inf"))
 
     if attn_mask is not None:
         if attn_mask.dtype == torch.bool:
-            attn_bias = attn_bias.masked_fill(attn_mask.logical_not(), float("-inf"))
+            # Preserve masked_fill semantics while using scalar choices so TensorRT
+            # can recognize the compact additive-attention-mask pattern.
+            zero = torch.full((), 0.0, dtype=query.dtype, device=device)
+            negative = torch.full((), float("-inf"), dtype=query.dtype, device=device)
+            attn_bias = torch.where(attn_mask.logical_not(), negative, zero)
         else:
+            attn_bias = torch.zeros((L, S), dtype=query.dtype, device=device)
             attn_bias = attn_mask + attn_bias
 
     if enable_gqa:

tests/py/dynamo/conversion/test_arange_aten.py

@@ -1,15 +1,47 @@
-import unittest
-
+import tensorrt as trt
 import torch
 import torch.nn as nn
 import torch_tensorrt
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion import impl
+from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
 
 from .harness import DispatchTestCase
 
 
 class TestArangeConverter(DispatchTestCase):
+    def test_static_arange_uses_linspace_fill(self):
+        logger = trt.Logger(trt.Logger.ERROR)
+        builder = trt.Builder(logger)
+        network = builder.create_network(
+            1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
+        )
+        ctx = ConversionContext(network)
+
+        output = impl.arange.arange(
+            ctx,
+            torch.ops.aten.arange.start_step,
+            SourceIR.ATEN,
+            "arange",
+            start=0,
+            end=5,
+            step=1,
+        )
+
+        fill_layers = [
+            network.get_layer(index)
+            for index in range(network.num_layers)
+            if network.get_layer(index).type == trt.LayerType.FILL
+        ]
+        self.assertEqual(len(fill_layers), 1)
+        self.assertEqual(
+            fill_layers[0].name,
+            "[FILL]-[aten_ops.arange.start_step]-[arange_arange_fill]",
+        )
+        self.assertEqual(tuple(output.shape), (5,))
+
     @parameterized.expand(
         [
             (0, 5, 1),

tests/py/dynamo/lowering/test_decompositions.py

@@ -1956,6 +1956,61 @@ def forward(self, query, key, value):
             )
             self.assertEqual(lowered.module()(*inputs).dtype, torch.float16)
 
+    def test_lowering_scaled_dot_product_attention_bool_mask_bias(self):
+        class TestModule(torch.nn.Module):
+            def forward(self, query, key, value, attn_mask):
+                return torch.ops.aten.scaled_dot_product_attention.default(
+                    query, key, value, attn_mask
+                )
+
+        inputs = (
+            torch.randn(1, 2, 8, 16, dtype=torch.bfloat16, device="cuda"),
+            torch.randn(1, 2, 8, 16, dtype=torch.bfloat16, device="cuda"),
+            torch.randn(1, 2, 8, 16, dtype=torch.bfloat16, device="cuda"),
+            torch.ones(1, 1, 8, 8, dtype=torch.bool, device="cuda").tril(),
+        )
+        exported_program = torch.export.export(TestModule(), inputs)
+        lowered = exported_program.run_decompositions(
+            get_decompositions(decompose_attention=True)
+        )
+
+        full_nodes = [
+            node
+            for node in lowered.graph.nodes
+            if node.op == "call_function" and node.target == torch.ops.aten.full.default
+        ]
+        self.assertEqual(len(full_nodes), 2)
+        self.assertTrue(all(node.args[0] == [] for node in full_nodes))
+        self.assertEqual(
+            {node.args[1] for node in full_nodes},
+            {0.0, float("-inf")},
+        )
+        self.assertTrue(
+            all(node.meta["val"].dtype == torch.bfloat16 for node in full_nodes)
+        )
+
+        where_nodes = [
+            node
+            for node in lowered.graph.nodes
+            if node.op == "call_function" and node.target == torch.ops.aten.where.self
+        ]
+        self.assertEqual(len(where_nodes), 1)
+        self.assertEqual(where_nodes[0].meta["val"].dtype, torch.bfloat16)
+        logical_not_nodes = [
+            node
+            for node in lowered.graph.nodes
+            if node.op == "call_function"
+            and node.target == torch.ops.aten.logical_not.default
+        ]
+        self.assertEqual(len(logical_not_nodes), 1)
+        self.assertIs(where_nodes[0].args[0], logical_not_nodes[0])
+        torch.testing.assert_close(
+            lowered.module()(*inputs),
+            exported_program.module()(*inputs),
+            rtol=RTOL,
+            atol=ATOL,
+        )
+
     @parameterized.expand(
         [
             (True, False, None, False),