MLA attention fallback

py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py

@@ -14,6 +14,7 @@
 from .complex_graph_rewrite import complex_graph_detection
 from .constant_folding import constant_fold
 from .decompose_dynamic_slice_scatter import decompose_dynamic_slice_scatter
+from .decompose_unsupported_attention import decompose_unsupported_attention
 from .eliminate_sym_min_int64_max import eliminate_sym_min_int64_max
 from .force_causal_efficient_attention import force_causal_efficient_attention
 from .fuse_prims_broadcast import fuse_prims_broadcast
@@ -45,6 +46,7 @@
     remove_num_users_is_0_nodes,
     complex_graph_detection,
     force_causal_efficient_attention,
+    decompose_unsupported_attention,
     eliminate_sym_min_int64_max,
     normalize_negative_slice_stop,
     annotate_fp8_sdpa,

py/torch_tensorrt/dynamo/lowering/passes/decompose_unsupported_attention.py

@@ -0,0 +1,130 @@
+import logging
+from typing import Any, Callable, Dict, Optional, Tuple
+
+import torch
+from torch.fx.experimental.proxy_tensor import make_fx
+from torch_tensorrt.dynamo._settings import CompilationSettings
+from torch_tensorrt.dynamo.conversion.aten_ops_converters import (
+    scaled_dot_product_attention_validator,
+)
+from torch_tensorrt.dynamo.lowering._decompositions import (
+    scaled_dot_product_attention_decomposition,
+)
+from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
+    clean_up_graph_after_modifications,
+)
+
+logger = logging.getLogger(__name__)
+
+# Kept out of the default decomp table so IAttention can consume it when valid.
+_ATTENTION_FALLBACKS: Dict[
+    Any,
+    Tuple[
+        Callable[[torch.fx.Node, Optional[CompilationSettings]], bool],
+        Callable[..., Any],
+    ],
+] = {
+    torch.ops.aten.scaled_dot_product_attention.default: (
+        scaled_dot_product_attention_validator,
+        scaled_dot_product_attention_decomposition,
+    ),
+}
+
+
+def _example_val(x: Any) -> Any:
+    if isinstance(x, torch.fx.Node):
+        val = x.meta.get("val")
+        if val is None:
+            raise RuntimeError(
+                f"Cannot decompose attention node: {x} has no meta['val']"
+            )
+        return val
+    return x
+
+
+def _inline_traced_decomp(
+    gm: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    decomp_fn: Callable[..., Any],
+    use_fp32_acc: bool,
+) -> None:
+    """Replace ``node`` with the aten subgraph from tracing ``decomp_fn``."""
+
+    example_args = tuple(_example_val(a) for a in node.args)
+    example_kwargs = {k: _example_val(v) for k, v in node.kwargs.items()}
+    example_kwargs["use_fp32_acc"] = use_fp32_acc
+
+    def wrapped(*args: Any) -> Any:
+        return decomp_fn(*args, **example_kwargs)
+
+    traced = make_fx(wrapped, tracing_mode="fake")(*example_args)
+
+    arg_nodes = list(node.args)
+    env: Dict[torch.fx.Node, Any] = {}
+    ph_idx = 0
+    output: Any = None
+
+    with gm.graph.inserting_before(node):
+        for n in traced.graph.nodes:
+            if n.op == "placeholder":
+                env[n] = arg_nodes[ph_idx]
+                ph_idx += 1
+            elif n.op == "get_attr":
+                const = getattr(traced, n.target)
+                const_name = f"_attn_fb_{node.name}_{n.target.replace('.', '_')}"
+                gm.register_buffer(const_name, const.detach().clone())
+                env[n] = gm.graph.get_attr(const_name)
+            elif n.op == "call_function":
+                new_args = torch.fx.node.map_arg(n.args, lambda x: env[x])
+                new_kwargs = torch.fx.node.map_arg(n.kwargs, lambda x: env[x])
+                new_n = gm.graph.call_function(n.target, new_args, new_kwargs)
+                if "val" in n.meta:
+                    new_n.meta["val"] = n.meta["val"]
+                env[n] = new_n
+            elif n.op == "output":
+                output = torch.fx.node.map_arg(n.args[0], lambda x: env[x])
+            else:
+                raise RuntimeError(
+                    f"Unexpected op {n.op} while inlining attention decomposition"
+                )
+
+    assert output is not None
+    node.replace_all_uses_with(output)
+    gm.graph.erase_node(node)
+
+
+def decompose_unsupported_attention(
+    gm: torch.fx.GraphModule, settings: CompilationSettings
+) -> torch.fx.GraphModule:
+    """Per-node fallback when the native attention converter declines a node.
+
+    Attention ops are deliberately left out of the default decomposition table so
+    ``IAttentionLayer`` can consume them whole. When the capability validator
+    rejects a node (e.g. MLA where K and V head dims differ), leave the rest of
+    the graph on the native path and decompose only the declined nodes.
+    """
+    if settings.decompose_attention:
+        return gm
+
+    changed = False
+    for node in list(gm.graph.nodes):
+        if node.op != "call_function" or node.target not in _ATTENTION_FALLBACKS:
+            continue
+
+        validator, decomp = _ATTENTION_FALLBACKS[node.target]
+        if validator(node, settings):
+            continue
+
+        logger.info(
+            "Attention converter declined %s; decomposing that node in place "
+            "(native IAttention kept for other attention nodes)",
+            node.name,
+        )
+        _inline_traced_decomp(gm, node, decomp, settings.use_fp32_acc)
+        changed = True
+
+    if changed:
+        gm = clean_up_graph_after_modifications(gm)
+        logger.debug("After decompose_unsupported_attention:\n%s", gm.graph)
+
+    return gm

tests/py/dynamo/lowering/test_decompose_unsupported_attention.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+from torch.testing._internal.common_utils import TestCase, run_tests
+from torch_tensorrt.dynamo._settings import CompilationSettings
+from torch_tensorrt.dynamo.conversion.aten_ops_converters import (
+    scaled_dot_product_attention_validator,
+)
+from torch_tensorrt.dynamo.lowering.passes.decompose_unsupported_attention import (
+    decompose_unsupported_attention,
+)
+
+from ..testing_utilities import lower_graph_testing
+
+
+class TestDecomposeUnsupportedAttention(TestCase):
+    def test_mla_kv_head_dim_mismatch_is_decomposed(self):
+        # MLA: K head dim = nope + rope, V head dim = v only.
+        b, h, s, d_k, d_v = 1, 2, 4, 6, 4
+
+        class MLA(nn.Module):
+            def forward(self, q, k, v):
+                return torch.ops.aten.scaled_dot_product_attention.default(q, k, v)
+
+        inputs = [
+            torch.randn(b, h, s, d_k),
+            torch.randn(b, h, s, d_k),
+            torch.randn(b, h, s, d_v),
+        ]
+        ep = torch.export.export(MLA(), tuple(inputs))
+        settings = CompilationSettings(min_block_size=1, decompose_attention=False)
+        gm = ep.module()
+
+        sdpa = next(
+            n
+            for n in gm.graph.nodes
+            if n.target == torch.ops.aten.scaled_dot_product_attention.default
+        )
+        self.assertFalse(scaled_dot_product_attention_validator(sdpa, settings))
+
+        gm = decompose_unsupported_attention(gm, settings)
+
+        targets = {n.target for n in gm.graph.nodes if n.op == "call_function"}
+        self.assertNotIn(torch.ops.aten.scaled_dot_product_attention.default, targets)
+
+    def test_equal_kv_shapes_left_intact(self):
+        class MHA(nn.Module):
+            def forward(self, q, k, v):
+                return torch.ops.aten.scaled_dot_product_attention.default(q, k, v)
+
+        inputs = [torch.randn(1, 2, 4, 8) for _ in range(3)]
+        ep = torch.export.export(MHA(), tuple(inputs))
+        settings = CompilationSettings(min_block_size=1, decompose_attention=False)
+        gm = ep.module()
+        gm = decompose_unsupported_attention(gm, settings)
+
+        targets = {n.target for n in gm.graph.nodes if n.op == "call_function"}
+        self.assertIn(torch.ops.aten.scaled_dot_product_attention.default, targets)
+
+    def test_post_lowering_mla_fully_supported(self):
+        class MLA(nn.Module):
+            def forward(self, q, k, v):
+                return torch.nn.functional.scaled_dot_product_attention(q, k, v)
+
+        inputs = [
+            torch.randn(1, 2, 4, 6, device="cuda"),
+            torch.randn(1, 2, 4, 6, device="cuda"),
+            torch.randn(1, 2, 4, 4, device="cuda"),
+        ]
+        unexpected_ops = {torch.ops.aten.scaled_dot_product_attention.default}
+        fx_graph = torch.export.export(MLA(), tuple(inputs)).module()
+        unexpected_ops_seen, _ = lower_graph_testing(
+            fx_graph,
+            inputs,
+            unexpected_ops=unexpected_ops,
+            min_block_size=1,
+            decompose_attention=False,
+        )
+        self.assertEqual(
+            len(unexpected_ops_seen),
+            0,
+            f"SDPA should have been decomposed for MLA shapes: {unexpected_ops_seen}",
+        )
+
+
+if __name__ == "__main__":
+    run_tests()