From 823324619bdf728f2f3dd44a1b776fdb8c9f2aea Mon Sep 17 00:00:00 2001 From: Yulong Wang <7679871+fs-eire@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:59:59 -0700 Subject: [PATCH 1/3] fix: align with ONNX graph in attention_decomposition --- .../dynamo/lowering/_decompositions.py | 13 +++-- .../py/dynamo/lowering/test_decompositions.py | 54 +++++++++++++++++++ 2 files changed, 63 insertions(+), 4 deletions(-) diff --git a/py/torch_tensorrt/dynamo/lowering/_decompositions.py b/py/torch_tensorrt/dynamo/lowering/_decompositions.py index a32aef1b28..e72557f436 100644 --- a/py/torch_tensorrt/dynamo/lowering/_decompositions.py +++ b/py/torch_tensorrt/dynamo/lowering/_decompositions.py @@ -485,18 +485,23 @@ 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")) + # Keep the mask condition in its valid=True orientation and use scalar + # constants so TensorRT can recognize the additive-attention-mask pattern. + zero = torch.full((), 0.0, dtype=query.dtype, device=device) + negative = torch.full( + (), torch.finfo(query.dtype).min, dtype=query.dtype, device=device + ) + attn_bias = torch.where(attn_mask, zero, negative) else: + attn_bias = torch.zeros((L, S), dtype=query.dtype, device=device) attn_bias = attn_mask + attn_bias if enable_gqa: diff --git a/tests/py/dynamo/lowering/test_decompositions.py b/tests/py/dynamo/lowering/test_decompositions.py index e1d7e87ea2..a40705cc6d 100644 --- a/tests/py/dynamo/lowering/test_decompositions.py +++ b/tests/py/dynamo/lowering/test_decompositions.py @@ -1956,6 +1956,60 @@ 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, torch.finfo(torch.bfloat16).min}, + ) + 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) + self.assertFalse( + any( + node.op == "call_function" + and node.target == torch.ops.aten.logical_not.default + for node in lowered.graph.nodes + ) + ) + torch.testing.assert_close( + lowered.module()(*inputs), + exported_program.module()(*inputs), + rtol=RTOL, + atol=ATOL, + ) + @parameterized.expand( [ (True, False, None, False), From 237c67147c0e3bd2b67b3b86eae0494d3aacbd1c Mon Sep 17 00:00:00 2001 From: Yulong Wang <7679871+fs-eire@users.noreply.github.com> Date: Wed, 22 Jul 2026 06:57:31 +0000 Subject: [PATCH 2/3] fix: limit attention mask change to scalar choices Restore the original -inf value and invalid-mask polarity while retaining scalar Select inputs required for compact TensorRT MHA masks. --- .../dynamo/lowering/_decompositions.py | 10 ++++------ tests/py/dynamo/lowering/test_decompositions.py | 17 +++++++++-------- 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/py/torch_tensorrt/dynamo/lowering/_decompositions.py b/py/torch_tensorrt/dynamo/lowering/_decompositions.py index e72557f436..defb279190 100644 --- a/py/torch_tensorrt/dynamo/lowering/_decompositions.py +++ b/py/torch_tensorrt/dynamo/lowering/_decompositions.py @@ -493,13 +493,11 @@ def scaled_dot_product_attention_decomposition( if attn_mask is not None: if attn_mask.dtype == torch.bool: - # Keep the mask condition in its valid=True orientation and use scalar - # constants so TensorRT can recognize the additive-attention-mask pattern. + # 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( - (), torch.finfo(query.dtype).min, dtype=query.dtype, device=device - ) - attn_bias = torch.where(attn_mask, zero, negative) + 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 diff --git a/tests/py/dynamo/lowering/test_decompositions.py b/tests/py/dynamo/lowering/test_decompositions.py index a40705cc6d..615e8161ce 100644 --- a/tests/py/dynamo/lowering/test_decompositions.py +++ b/tests/py/dynamo/lowering/test_decompositions.py @@ -1983,7 +1983,7 @@ def forward(self, query, key, value, attn_mask): self.assertTrue(all(node.args[0] == [] for node in full_nodes)) self.assertEqual( {node.args[1] for node in full_nodes}, - {0.0, torch.finfo(torch.bfloat16).min}, + {0.0, float("-inf")}, ) self.assertTrue( all(node.meta["val"].dtype == torch.bfloat16 for node in full_nodes) @@ -1996,13 +1996,14 @@ def forward(self, query, key, value, attn_mask): ] self.assertEqual(len(where_nodes), 1) self.assertEqual(where_nodes[0].meta["val"].dtype, torch.bfloat16) - self.assertFalse( - any( - node.op == "call_function" - and node.target == torch.ops.aten.logical_not.default - for node in lowered.graph.nodes - ) - ) + 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), From 7a7c41ace69bc3bc348cdc62b72ce27835c7067d Mon Sep 17 00:00:00 2001 From: Yulong Wang <7679871+fs-eire@users.noreply.github.com> Date: Wed, 22 Jul 2026 06:47:37 +0000 Subject: [PATCH 3/3] fix: preserve static arange as TensorRT Fill --- .../dynamo/conversion/impl/arange.py | 33 +++++++++++------ .../py/dynamo/conversion/test_arange_aten.py | 36 +++++++++++++++++-- 2 files changed, 56 insertions(+), 13 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/arange.py b/py/torch_tensorrt/dynamo/conversion/impl/arange.py index f8db7e803f..ac42bc90a4 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/arange.py +++ b/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) diff --git a/tests/py/dynamo/conversion/test_arange_aten.py b/tests/py/dynamo/conversion/test_arange_aten.py index 7705590e2e..46e808b81a 100644 --- a/tests/py/dynamo/conversion/test_arange_aten.py +++ b/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),