fix: skip const-fold for arange for attention mask

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,11 @@ 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. 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,
@@ -357,6 +363,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 +489,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 +584,11 @@ 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. 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,
@@ -765,6 +778,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 +1729,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 +1828,11 @@ 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. 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:
@@ -1978,6 +1998,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,
     }
 

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
 

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,11 @@ 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. 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.
     """
@@ -178,9 +198,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 +225,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 +256,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",
 }
 

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:

py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py

@@ -0,0 +1,15 @@
+from ._core import (
+    CONSTANT_FOLD_EXCLUSION_META_KEY,
+    ConstantFoldExclusionRule,
+    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",
+    "validate_disabled_constant_fold_exclusions",
+]

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

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)

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,