From 7c180314c668742d9d6003ac8a5c03a16fb12d3d Mon Sep 17 00:00:00 2001 From: Yulong Wang <7679871+fs-eire@users.noreply.github.com> Date: Thu, 30 Jul 2026 11:38:41 -0700 Subject: [PATCH 1/2] feat: support configurable constant-fold exclusions --- py/torch_tensorrt/dynamo/_compiler.py | 18 ++++ py/torch_tensorrt/dynamo/_defaults.py | 1 + py/torch_tensorrt/dynamo/_settings.py | 35 ++++++++ .../constant_fold_exclusions/__init__.py | 13 +++ .../constant_fold_exclusions/_core.py | 61 +++++++++++++ .../lowering/passes/_aten_lowering_pass.py | 4 + .../lowering/passes/constant_folding.py | 7 +- .../passes/mark_constant_fold_exclusions.py | 40 +++++++++ .../lowering/test_constant_fold_exclusions.py | 85 +++++++++++++++++++ 9 files changed, 263 insertions(+), 1 deletion(-) create mode 100644 py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py create mode 100644 py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py create mode 100644 py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py create mode 100644 tests/py/dynamo/lowering/test_constant_fold_exclusions.py diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index 3ebac2a21f..b6a1321342 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -117,6 +117,7 @@ def cross_compile_for_windows( cpu_memory_budget: Optional[int] = _defaults.CPU_MEMORY_BUDGET, dynamically_allocate_resources: bool = _defaults.DYNAMICALLY_ALLOCATE_RESOURCES, decompose_attention: bool = _defaults.DECOMPOSE_ATTENTION, + disabled_constant_fold_exclusions: Collection[str] = _defaults.DISABLED_CONSTANT_FOLD_EXCLUSIONS, attn_bias_is_causal: bool = _defaults.ATTN_BIAS_IS_CAUSAL, fallback_data_dependent_ops: bool = _defaults.FALLBACK_DATA_DEPENDENT_OPS, **kwargs: Any, @@ -201,6 +202,10 @@ def cross_compile_for_windows( instead of using the attention converters. When combined with ``use_fp32_acc=True``, decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only the final output back to FP16. + disabled_constant_fold_exclusions (Collection[str]): IDs of + Torch-TensorRT rules that exclude matching FX nodes from constant + folding. Naming a rule here turns it off, so the nodes it would have + kept become foldable again. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. **kwargs: Any, @@ -357,6 +362,7 @@ def cross_compile_for_windows( "cpu_memory_budget": cpu_memory_budget, "dynamically_allocate_resources": dynamically_allocate_resources, "decompose_attention": decompose_attention, + "disabled_constant_fold_exclusions": disabled_constant_fold_exclusions, "attn_bias_is_causal": attn_bias_is_causal, "fallback_data_dependent_ops": fallback_data_dependent_ops, } @@ -482,6 +488,7 @@ def compile( enable_resource_partitioning: bool = _defaults.ENABLE_RESOURCE_PARTITIONING, dynamically_allocate_resources: bool = _defaults.DYNAMICALLY_ALLOCATE_RESOURCES, decompose_attention: bool = _defaults.DECOMPOSE_ATTENTION, + disabled_constant_fold_exclusions: Collection[str] = _defaults.DISABLED_CONSTANT_FOLD_EXCLUSIONS, attn_bias_is_causal: bool = _defaults.ATTN_BIAS_IS_CAUSAL, fallback_data_dependent_ops: bool = _defaults.FALLBACK_DATA_DEPENDENT_OPS, **kwargs: Any, @@ -576,6 +583,10 @@ def compile( instead of using the attention converters. When combined with ``use_fp32_acc=True``, decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only the final output back to FP16. + disabled_constant_fold_exclusions (Collection[str]): IDs of + Torch-TensorRT rules that exclude matching FX nodes from constant + folding. Naming a rule here turns it off, so the nodes it would have + kept become foldable again. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. **kwargs: Any, @@ -765,6 +776,7 @@ def compile( "cpu_memory_budget": cpu_memory_budget, "dynamically_allocate_resources": dynamically_allocate_resources, "decompose_attention": decompose_attention, + "disabled_constant_fold_exclusions": disabled_constant_fold_exclusions, "attn_bias_is_causal": attn_bias_is_causal, "fallback_data_dependent_ops": fallback_data_dependent_ops, } @@ -1715,6 +1727,7 @@ def convert_exported_program_to_serialized_trt_engine( offload_module_to_cpu: bool = _defaults.OFFLOAD_MODULE_TO_CPU, use_distributed_mode_trace: bool = _defaults.USE_DISTRIBUTED_MODE_TRACE, decompose_attention: bool = _defaults.DECOMPOSE_ATTENTION, + disabled_constant_fold_exclusions: Collection[str] = _defaults.DISABLED_CONSTANT_FOLD_EXCLUSIONS, attn_bias_is_causal: bool = _defaults.ATTN_BIAS_IS_CAUSAL, lift_mutable_buffers: bool = False, arg_input_binding_names: Any = None, @@ -1813,6 +1826,10 @@ def convert_exported_program_to_serialized_trt_engine( instead of using the attention converters. When combined with ``use_fp32_acc=True``, decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only the final output back to FP16. + disabled_constant_fold_exclusions (Collection[str]): IDs of + Torch-TensorRT rules that exclude matching FX nodes from constant + folding. Naming a rule here turns it off, so the nodes it would have + kept become foldable again. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. **kwargs: Any, Returns: @@ -1978,6 +1995,7 @@ def convert_exported_program_to_serialized_trt_engine( "offload_module_to_cpu": offload_module_to_cpu, "use_distributed_mode_trace": use_distributed_mode_trace, "decompose_attention": decompose_attention, + "disabled_constant_fold_exclusions": disabled_constant_fold_exclusions, "attn_bias_is_causal": attn_bias_is_causal, } diff --git a/py/torch_tensorrt/dynamo/_defaults.py b/py/torch_tensorrt/dynamo/_defaults.py index 9c8a1f9f90..ce6295cf66 100644 --- a/py/torch_tensorrt/dynamo/_defaults.py +++ b/py/torch_tensorrt/dynamo/_defaults.py @@ -65,6 +65,7 @@ CPU_MEMORY_BUDGET = None DYNAMICALLY_ALLOCATE_RESOURCES = False DECOMPOSE_ATTENTION = False +DISABLED_CONSTANT_FOLD_EXCLUSIONS = frozenset[str]() ATTN_BIAS_IS_CAUSAL = True FALLBACK_DATA_DEPENDENT_OPS = False diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py index 7ff5f6bdff..243c9bb94b 100644 --- a/py/torch_tensorrt/dynamo/_settings.py +++ b/py/torch_tensorrt/dynamo/_settings.py @@ -18,6 +18,7 @@ CACHE_BUILT_ENGINES, CPU_MEMORY_BUDGET, DECOMPOSE_ATTENTION, + DISABLED_CONSTANT_FOLD_EXCLUSIONS, DISABLE_TF32, DLA_GLOBAL_DRAM_SIZE, DLA_LOCAL_DRAM_SIZE, @@ -58,6 +59,20 @@ ) +def _normalize_disabled_constant_fold_exclusions( + rule_ids: Collection[str], +) -> Set[str]: + # Deferred import: torch_tensorrt.dynamo.lowering.passes imports this module, + # so importing the rule registry at module scope would be circular. Reaching + # for the core directly is safe from any point of that cycle because it + # depends on nothing but torch. + from torch_tensorrt.dynamo.lowering.constant_fold_exclusions._core import ( + validate_disabled_constant_fold_exclusions, + ) + + return validate_disabled_constant_fold_exclusions(rule_ids) + + @dataclass class CompilationSettings: """Compilation settings for Torch-TensorRT Dynamo Paths @@ -119,6 +134,10 @@ class CompilationSettings: instead of using the attention converters. When combined with ``use_fp32_acc=True``, decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only the final output back to FP16. + disabled_constant_fold_exclusions (Collection[str]): IDs of Torch-TensorRT + rules that exclude matching FX nodes from constant folding. Naming a + rule here turns it off, so the nodes it would have kept become + foldable again. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. """ @@ -178,9 +197,19 @@ class CompilationSettings: cpu_memory_budget: Optional[int] = CPU_MEMORY_BUDGET dynamically_allocate_resources: bool = DYNAMICALLY_ALLOCATE_RESOURCES decompose_attention: bool = DECOMPOSE_ATTENTION + disabled_constant_fold_exclusions: Collection[str] = ( + DISABLED_CONSTANT_FOLD_EXCLUSIONS + ) attn_bias_is_causal: bool = ATTN_BIAS_IS_CAUSAL fallback_data_dependent_ops: bool = FALLBACK_DATA_DEPENDENT_OPS + def __post_init__(self) -> None: + self.disabled_constant_fold_exclusions = ( + _normalize_disabled_constant_fold_exclusions( + self.disabled_constant_fold_exclusions + ) + ) + def __getstate__(self) -> dict[str, Any]: from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( ConverterRegistry, @@ -195,6 +224,11 @@ def __getstate__(self) -> dict[str, Any]: def __setstate__(self, state: dict[str, Any]) -> None: state.pop("use_python_runtime", None) + state["disabled_constant_fold_exclusions"] = ( + _normalize_disabled_constant_fold_exclusions( + state.get("disabled_constant_fold_exclusions", ()) + ) + ) state.setdefault("fallback_data_dependent_ops", FALLBACK_DATA_DEPENDENT_OPS) self.__dict__.update(state) @@ -221,6 +255,7 @@ def __setstate__(self, state: dict[str, Any]) -> None: "autocast_max_depth_of_reduction", "autocast_calibration_dataloader", "decompose_attention", + "disabled_constant_fold_exclusions", "attn_bias_is_causal", } diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py new file mode 100644 index 0000000000..c591d312ab --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py @@ -0,0 +1,13 @@ +from ._core import ( + CONSTANT_FOLD_EXCLUSION_META_KEY, + ConstantFoldExclusionRule, + register_constant_fold_exclusion_rule, + validate_disabled_constant_fold_exclusions, +) + +__all__ = [ + "CONSTANT_FOLD_EXCLUSION_META_KEY", + "ConstantFoldExclusionRule", + "register_constant_fold_exclusion_rule", + "validate_disabled_constant_fold_exclusions", +] diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py new file mode 100644 index 0000000000..7d38f385c6 --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py @@ -0,0 +1,61 @@ +from typing import Callable, Collection, Iterable + +import torch + +CONSTANT_FOLD_EXCLUSION_META_KEY = "_torch_tensorrt_constant_fold_exclusions" + +ConstantFoldExclusionRule = Callable[[torch.fx.Node], Iterable[torch.fx.Node]] +_CONSTANT_FOLD_EXCLUSION_RULES: dict[str, ConstantFoldExclusionRule] = {} + + +def _mark_constant_fold_exclusion(nodes: Iterable[torch.fx.Node], rule_id: str) -> None: + """Record which rule wants ``nodes`` kept out of constant folding. + + The marks carry their rule ID rather than a bare flag so the exclusion pass + can revoke the ones belonging to disabled rules, whichever marking path + produced them. + """ + for node in nodes: + node.meta.setdefault(CONSTANT_FOLD_EXCLUSION_META_KEY, set()).add(rule_id) + + +def register_constant_fold_exclusion_rule( + rule_id: str, +) -> Callable[[ConstantFoldExclusionRule], ConstantFoldExclusionRule]: + """Register a named rule that selects FX nodes to exclude from folding.""" + if not isinstance(rule_id, str) or not rule_id: + raise ValueError("A constant-fold exclusion rule ID must be a non-empty string") + + def register(rule: ConstantFoldExclusionRule) -> ConstantFoldExclusionRule: + if rule_id in _CONSTANT_FOLD_EXCLUSION_RULES: + raise ValueError( + f"Constant-fold exclusion rule {rule_id!r} is already registered" + ) + + _CONSTANT_FOLD_EXCLUSION_RULES[rule_id] = rule + return rule + + return register + + +def validate_disabled_constant_fold_exclusions( + rule_ids: Collection[str], +) -> set[str]: + """Validate disabled rule IDs and return them as a set.""" + if isinstance(rule_ids, str): + raise TypeError( + "disabled_constant_fold_exclusions must be a collection of rule IDs, " + "not a single string" + ) + + disabled_rule_ids = set(rule_ids) + unknown_rule_ids = disabled_rule_ids - _CONSTANT_FOLD_EXCLUSION_RULES.keys() + if unknown_rule_ids: + available_rule_ids = ", ".join(sorted(_CONSTANT_FOLD_EXCLUSION_RULES)) + raise ValueError( + "Unknown constant-fold exclusion rule IDs: " + f"{sorted(unknown_rule_ids)}. Available rule IDs: " + f"[{available_rule_ids}]" + ) + + return disabled_rule_ids diff --git a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py index 4f41d3e6b7..d141feedd1 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py @@ -6,6 +6,9 @@ from torch_tensorrt._utils import is_tegra_platform from torch_tensorrt.dynamo._settings import CompilationSettings from torch_tensorrt.dynamo.lowering.passes._FakeTensorUpdater import FakeTensorUpdater +from torch_tensorrt.dynamo.lowering.passes.mark_constant_fold_exclusions import ( + mark_constant_fold_exclusions, +) from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( trace_intermediate_node_outputs, ) @@ -37,6 +40,7 @@ post_lowering_pass_list = [ replace_fused_rms_norm, remove_input_alias_fixing_clones, + mark_constant_fold_exclusions, constant_fold, repair_input_as_output, fuse_prims_broadcast, diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py index 60dec56f3b..a88aed4a9c 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py @@ -4,6 +4,9 @@ import torch from torch_tensorrt._utils import sanitized_torch_version from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.lowering.constant_fold_exclusions import ( + CONSTANT_FOLD_EXCLUSION_META_KEY, +) from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( clean_up_graph_after_modifications, ) @@ -125,4 +128,6 @@ def is_impure(self, node: torch.fx.node.Node) -> bool: if node.target in self.quantization_ops: return True - return False + # The meta value holds the IDs of the rules that marked this node; it is + # empty once every rule that claimed it has been disabled. + return bool(node.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY)) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py b/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py new file mode 100644 index 0000000000..a6856680e6 --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py @@ -0,0 +1,40 @@ +from typing import Any, Optional + +import torch +from torch_tensorrt.dynamo.lowering.constant_fold_exclusions._core import ( + CONSTANT_FOLD_EXCLUSION_META_KEY, + _CONSTANT_FOLD_EXCLUSION_RULES, + _mark_constant_fold_exclusion, + validate_disabled_constant_fold_exclusions, +) + + +def mark_constant_fold_exclusions( + gm: torch.fx.GraphModule, settings: Optional[Any] = None +) -> torch.fx.GraphModule: + """Apply the registered rules that exclude FX nodes from constant folding. + + This pass is the single authority on which rules are in effect. It runs + immediately before ``constant_fold`` and is the only marking path that sees + ``settings``: rules that mark nodes while a decomposition is traced run + during ``run_decompositions``, long before a settings object is reachable. + Those marks are therefore revoked here rather than suppressed where they are + made, so a caller only has to communicate the disabled rules once. + """ + disabled_rule_ids = validate_disabled_constant_fold_exclusions( + settings.disabled_constant_fold_exclusions if settings is not None else () + ) + + for node in gm.graph.nodes: + for rule_id, rule in _CONSTANT_FOLD_EXCLUSION_RULES.items(): + if rule_id in disabled_rule_ids: + continue + _mark_constant_fold_exclusion(rule(node), rule_id) + + if disabled_rule_ids: + for node in gm.graph.nodes: + marking_rule_ids = node.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY) + if marking_rule_ids: + marking_rule_ids -= disabled_rule_ids + + return gm diff --git a/tests/py/dynamo/lowering/test_constant_fold_exclusions.py b/tests/py/dynamo/lowering/test_constant_fold_exclusions.py new file mode 100644 index 0000000000..b421003d83 --- /dev/null +++ b/tests/py/dynamo/lowering/test_constant_fold_exclusions.py @@ -0,0 +1,85 @@ +import torch +from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt.dynamo._settings import ( + CompilationSettings, + settings_are_compatible, +) +from torch_tensorrt.dynamo.lowering.constant_fold_exclusions import ( + CONSTANT_FOLD_EXCLUSION_META_KEY, + register_constant_fold_exclusion_rule, +) +from torch_tensorrt.dynamo.lowering.passes.mark_constant_fold_exclusions import ( + mark_constant_fold_exclusions, +) + +TEST_RULE_ID = "test_arbitrary_node" + + +def custom_target(): + return torch.ones(1) + + +@register_constant_fold_exclusion_rule(TEST_RULE_ID) +def custom_rule(node): + return (node,) if node.target is custom_target else () + + +class TestConstantFoldExclusionMechanics(TestCase): + def _custom_graph(self): + graph = torch.fx.Graph() + custom_node = graph.call_function(custom_target) + graph.output(custom_node) + return torch.fx.GraphModule({}, graph), custom_node + + def test_disabled_rules_default_to_empty(self): + self.assertEqual(CompilationSettings().disabled_constant_fold_exclusions, set()) + self.assertEqual( + CompilationSettings( + disabled_constant_fold_exclusions=[TEST_RULE_ID] + ).disabled_constant_fold_exclusions, + {TEST_RULE_ID}, + ) + with self.assertRaisesRegex(TypeError, "collection of rule IDs"): + CompilationSettings(disabled_constant_fold_exclusions=TEST_RULE_ID) + + def test_old_serialized_setting_defaults_to_no_disabled_rules(self): + state = CompilationSettings().__dict__.copy() + state.pop("disabled_constant_fold_exclusions") + restored = CompilationSettings.__new__(CompilationSettings) + restored.__setstate__(state) + self.assertEqual(restored.disabled_constant_fold_exclusions, set()) + + def test_setting_changes_engine_compatibility(self): + compatible, incompatible_settings = settings_are_compatible( + CompilationSettings(), + CompilationSettings(disabled_constant_fold_exclusions={TEST_RULE_ID}), + ) + self.assertFalse(compatible) + self.assertIn( + "disabled_constant_fold_exclusions", + incompatible_settings, + ) + + def test_registered_rule_can_mark_an_arbitrary_node(self): + gm, custom_node = self._custom_graph() + mark_constant_fold_exclusions(gm) + self.assertTrue(custom_node.meta[CONSTANT_FOLD_EXCLUSION_META_KEY]) + + def test_disabled_rule_does_not_mark_a_node(self): + gm, custom_node = self._custom_graph() + mark_constant_fold_exclusions( + gm, + CompilationSettings(disabled_constant_fold_exclusions={TEST_RULE_ID}), + ) + self.assertFalse(custom_node.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY, False)) + + def test_unknown_disabled_rule_is_rejected(self): + with self.assertRaisesRegex( + ValueError, + "Unknown constant-fold exclusion rule IDs", + ): + CompilationSettings(disabled_constant_fold_exclusions={"unknown_rule"}) + + +if __name__ == "__main__": + run_tests() From c2647ffa25d23de950e24e04cfd49d0b9465be1b Mon Sep 17 00:00:00 2001 From: Yulong Wang <7679871+fs-eire@users.noreply.github.com> Date: Thu, 30 Jul 2026 11:42:55 -0700 Subject: [PATCH 2/2] fix: preserve SDPA attention-mask aranges during constant folding --- py/torch_tensorrt/dynamo/_compiler.py | 9 +- py/torch_tensorrt/dynamo/_settings.py | 3 +- .../dynamo/lowering/_decompositions.py | 6 + .../constant_fold_exclusions/__init__.py | 2 + .../attention_mask.py | 106 ++++++++++++ ..._attention_mask_constant_fold_exclusion.py | 158 ++++++++++++++++++ 6 files changed, 280 insertions(+), 4 deletions(-) create mode 100644 py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py create mode 100644 tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index b6a1321342..b6a28dd7e3 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -205,7 +205,8 @@ def cross_compile_for_windows( disabled_constant_fold_exclusions (Collection[str]): IDs of Torch-TensorRT rules that exclude matching FX nodes from constant folding. Naming a rule here turns it off, so the nodes it would have - kept become foldable again. Default is empty. + kept become foldable again. Supported IDs: + ``"attention_mask_arange"``. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. **kwargs: Any, @@ -586,7 +587,8 @@ def compile( disabled_constant_fold_exclusions (Collection[str]): IDs of Torch-TensorRT rules that exclude matching FX nodes from constant folding. Naming a rule here turns it off, so the nodes it would have - kept become foldable again. Default is empty. + kept become foldable again. Supported IDs: + ``"attention_mask_arange"``. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. **kwargs: Any, @@ -1829,7 +1831,8 @@ def convert_exported_program_to_serialized_trt_engine( disabled_constant_fold_exclusions (Collection[str]): IDs of Torch-TensorRT rules that exclude matching FX nodes from constant folding. Naming a rule here turns it off, so the nodes it would have - kept become foldable again. Default is empty. + kept become foldable again. Supported IDs: + ``"attention_mask_arange"``. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. **kwargs: Any, Returns: diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py index 243c9bb94b..ece9dfa9b9 100644 --- a/py/torch_tensorrt/dynamo/_settings.py +++ b/py/torch_tensorrt/dynamo/_settings.py @@ -137,7 +137,8 @@ class CompilationSettings: disabled_constant_fold_exclusions (Collection[str]): IDs of Torch-TensorRT rules that exclude matching FX nodes from constant folding. Naming a rule here turns it off, so the nodes it would have kept become - foldable again. Default is empty. + foldable again. Supported IDs: ``"attention_mask_arange"``. + Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. """ diff --git a/py/torch_tensorrt/dynamo/lowering/_decompositions.py b/py/torch_tensorrt/dynamo/lowering/_decompositions.py index a32aef1b28..819b78f14f 100644 --- a/py/torch_tensorrt/dynamo/lowering/_decompositions.py +++ b/py/torch_tensorrt/dynamo/lowering/_decompositions.py @@ -14,6 +14,9 @@ from torch_tensorrt.dynamo.conversion.converter_utils import get_positive_dim from torch_tensorrt.dynamo.utils import to_torch_device +from .constant_fold_exclusions.attention_mask import ( + exclude_attn_mask_aranges_from_constant_fold, +) from ._decomposition_groups import ( ENABLED_TORCH_DECOMPOSITIONS, TORCH_TRT_DECOMPOSITIONS, @@ -494,6 +497,9 @@ def scaled_dot_product_attention_decomposition( attn_bias = attn_bias.masked_fill(temp_mask.logical_not(), float("-inf")) if attn_mask is not None: + # Unconditional: mark_constant_fold_exclusions revokes these marks when the + # rule is disabled, so the decompositions do not need the setting. + exclude_attn_mask_aranges_from_constant_fold(attn_mask) if attn_mask.dtype == torch.bool: attn_bias = attn_bias.masked_fill(attn_mask.logical_not(), float("-inf")) else: diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py index c591d312ab..d91137ac17 100644 --- a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py +++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py @@ -4,8 +4,10 @@ register_constant_fold_exclusion_rule, validate_disabled_constant_fold_exclusions, ) +from .attention_mask import ATTENTION_MASK_ARANGE_RULE_ID __all__ = [ + "ATTENTION_MASK_ARANGE_RULE_ID", "CONSTANT_FOLD_EXCLUSION_META_KEY", "ConstantFoldExclusionRule", "register_constant_fold_exclusion_rule", diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py new file mode 100644 index 0000000000..110ed3c4ab --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py @@ -0,0 +1,106 @@ +from typing import Any, Callable, Iterable + +import torch +from torch._subclasses.functional_tensor import mb_unwrap_functional_tensor +from torch.fx.experimental.proxy_tensor import get_proxy_mode, get_proxy_slot + +from ._core import ( + _mark_constant_fold_exclusion, + register_constant_fold_exclusion_rule, +) + +ATTENTION_MASK_ARANGE_RULE_ID = "attention_mask_arange" + + +def _is_arange_node(node: torch.fx.Node) -> bool: + return ( + node.op == "call_function" + and getattr(node.target, "overloadpacket", None) is torch.ops.aten.arange + ) + + +def _find_ancestor_nodes( + node: torch.fx.Node, + predicate: Callable[[torch.fx.Node], bool], +) -> list[torch.fx.Node]: + """Find ancestors of ``node`` that satisfy ``predicate``.""" + nodes_to_visit: list[torch.fx.Node] = [node] + visited: set[torch.fx.Node] = set() + matching_nodes: list[torch.fx.Node] = [] + + while nodes_to_visit: + current = nodes_to_visit.pop() + if current in visited: + continue + + visited.add(current) + if predicate(current): + matching_nodes.append(current) + + nodes_to_visit.extend(current.all_input_nodes) + + return matching_nodes + + +def _find_attention_mask_aranges( + mask_node: torch.fx.Node, +) -> list[torch.fx.Node]: + """Find aranges that contribute to an attention mask.""" + return _find_ancestor_nodes(mask_node, predicate=_is_arange_node) + + +def exclude_attn_mask_aranges_from_constant_fold(attn_mask: torch.Tensor) -> None: + """Mark aranges behind an attention mask while tracing a decomposition. + + ``run_decompositions`` invokes decompositions once for functionalization and + again while proxy tracing. Only the proxy-tracing invocation has an FX graph + on which metadata can be set. + """ + proxy_mode = get_proxy_mode() + if proxy_mode is None: + return + + unwrapped_mask = mb_unwrap_functional_tensor(attn_mask) + tracked_mask = get_proxy_slot( + unwrapped_mask, + proxy_mode.tracer, + default=None, + ) + proxy = getattr(tracked_mask, "proxy", tracked_mask) + mask_node = getattr(proxy, "node", None) + + if isinstance(mask_node, torch.fx.Node): + _mark_constant_fold_exclusion( + _find_attention_mask_aranges(mask_node), + ATTENTION_MASK_ARANGE_RULE_ID, + ) + + +@register_constant_fold_exclusion_rule(ATTENTION_MASK_ARANGE_RULE_ID) +def _attention_mask_arange_rule(node: torch.fx.Node) -> Iterable[torch.fx.Node]: + """Select aranges feeding an attention mask.""" + # Every SDPA overload that takes a mask keeps it at positional index 3. + # _scaled_dot_product_flash_attention is absent because it has no mask. + attention_mask_args: dict[Any, tuple[int, str]] = { + torch.ops.aten.scaled_dot_product_attention: (3, "attn_mask"), + torch.ops.aten._scaled_dot_product_efficient_attention: (3, "attn_bias"), + torch.ops.aten._scaled_dot_product_cudnn_attention: (3, "attn_bias"), + } + + if node.op != "call_function": + return () + + overload_packet = getattr(node.target, "overloadpacket", None) + mask_arg = attention_mask_args.get(overload_packet) + if mask_arg is None: + return () + + mask_index, mask_kwarg = mask_arg + mask = node.kwargs.get( + mask_kwarg, + node.args[mask_index] if len(node.args) > mask_index else None, + ) + if not isinstance(mask, torch.fx.Node): + return () + + return _find_attention_mask_aranges(mask) diff --git a/tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py b/tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py new file mode 100644 index 0000000000..8ce2b8efcc --- /dev/null +++ b/tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py @@ -0,0 +1,158 @@ +import torch +from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.lowering import ( + get_decompositions, + post_lowering, +) +from torch_tensorrt.dynamo.lowering.constant_fold_exclusions import ( + ATTENTION_MASK_ARANGE_RULE_ID, + CONSTANT_FOLD_EXCLUSION_META_KEY, +) +from torch_tensorrt.dynamo.lowering.passes.mark_constant_fold_exclusions import ( + mark_constant_fold_exclusions, +) + + +class TestAttentionMaskConstantFoldExclusion(TestCase): + class AttentionWithCausalMask(torch.nn.Module): + def forward(self, query, key, value, attention_mask): + sequence_length = query.shape[-2] + row = torch.arange(sequence_length, device=query.device) + col = torch.arange(sequence_length, device=query.device) + causal_mask = col.unsqueeze(0) <= row.unsqueeze(1) + combined_mask = causal_mask & attention_mask + unrelated_arange = torch.arange(3, device=query.device) + attention = torch.ops.aten.scaled_dot_product_attention.default( + query, + key, + value, + combined_mask, + ) + return attention, unrelated_arange + + def _export(self): + inputs = ( + torch.randn(1, 2, 8, 16, device="cuda"), + torch.randn(1, 2, 8, 16, device="cuda"), + torch.randn(1, 2, 8, 16, device="cuda"), + torch.ones(8, 8, dtype=torch.bool, device="cuda"), + ) + return torch.export.export(self.AttentionWithCausalMask(), inputs) + + def _assert_only_attention_aranges_survive( + self, decompose_attention, disabled_constant_fold_exclusions=() + ): + exported_program = self._export().run_decompositions( + get_decompositions(decompose_attention=decompose_attention) + ) + gm = post_lowering( + exported_program.module(), + CompilationSettings( + disabled_constant_fold_exclusions=disabled_constant_fold_exclusions + ), + ) + + arange_nodes = [ + node + for node in gm.graph.nodes + if node.op == "call_function" + and getattr(node.target, "overloadpacket", None) is torch.ops.aten.arange + ] + self.assertEqual(len(arange_nodes), 2) + self.assertTrue( + all( + node.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY, False) + for node in arange_nodes + ) + ) + + def test_decomposed_attention_mask_aranges_are_not_folded(self): + self._assert_only_attention_aranges_survive(decompose_attention=True) + + def test_ia_attention_mask_aranges_are_not_folded(self): + self._assert_only_attention_aranges_survive(decompose_attention=False) + + def test_native_attention_rules_can_be_disabled(self): + exported_program = self._export().run_decompositions( + get_decompositions(decompose_attention=False) + ) + gm = post_lowering( + exported_program.module(), + CompilationSettings( + disabled_constant_fold_exclusions={ATTENTION_MASK_ARANGE_RULE_ID} + ), + ) + self.assertFalse( + any( + node.op == "call_function" + and getattr(node.target, "overloadpacket", None) + is torch.ops.aten.arange + for node in gm.graph.nodes + ) + ) + + def test_decomposed_attention_rules_can_be_disabled(self): + """post_lowering is the only place a rule has to be disabled. + + The decompositions mark unconditionally while tracing, well before any + settings object is reachable, so post_lowering revokes those marks + instead of the caller having to communicate the setting twice. + """ + exported_program = self._export().run_decompositions( + get_decompositions(decompose_attention=True) + ) + gm = post_lowering( + exported_program.module(), + CompilationSettings( + disabled_constant_fold_exclusions={ATTENTION_MASK_ARANGE_RULE_ID} + ), + ) + self.assertFalse( + any( + node.op == "call_function" + and getattr(node.target, "overloadpacket", None) + is torch.ops.aten.arange + for node in gm.graph.nodes + ) + ) + + +class TestAttentionMaskArangeRuleCoverage(TestCase): + """Check every SDPA overload that carries an attention mask.""" + + MASKED_ATTENTION_OPS = ( + (torch.ops.aten.scaled_dot_product_attention.default, "attn_mask"), + (torch.ops.aten._scaled_dot_product_efficient_attention.default, "attn_bias"), + (torch.ops.aten._scaled_dot_product_cudnn_attention.default, "attn_bias"), + ) + + def _attention_graph(self, target, mask_kwarg): + graph = torch.fx.Graph() + query = graph.placeholder("query") + key = graph.placeholder("key") + value = graph.placeholder("value") + arange = graph.call_function(torch.ops.aten.arange.default, (8,)) + mask = graph.call_function(torch.ops.aten.unsqueeze.default, (arange, 0)) + if mask_kwarg is None: + attention = graph.call_function(target, (query, key, value, mask)) + else: + attention = graph.call_function( + target, (query, key, value), {mask_kwarg: mask} + ) + graph.output(attention) + return torch.fx.GraphModule({}, graph), arange + + def test_mask_aranges_are_marked_for_every_masked_attention_op(self): + for target, mask_kwarg in self.MASKED_ATTENTION_OPS: + for passed_as_kwarg in (False, True): + with self.subTest(target=target, passed_as_kwarg=passed_as_kwarg): + gm, arange = self._attention_graph( + target, mask_kwarg if passed_as_kwarg else None + ) + mark_constant_fold_exclusions(gm) + self.assertTrue(arange.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY)) + + +if __name__ == "__main__": + run_tests()