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 b92165f83f..4c24d59218 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py @@ -17,6 +17,7 @@ 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 +from .lower_associative_scan import lower_associative_scan from .normalize_negative_slice_stop import normalize_negative_slice_stop from .pass_manager import DynamoPassManager from .remove_assert_nodes import remove_assert_nodes @@ -37,6 +38,7 @@ post_lowering_pass_list = [ replace_fused_rms_norm, remove_input_alias_fixing_clones, + lower_associative_scan, constant_fold, repair_input_as_output, fuse_prims_broadcast, diff --git a/py/torch_tensorrt/dynamo/lowering/passes/lower_associative_scan.py b/py/torch_tensorrt/dynamo/lowering/passes/lower_associative_scan.py new file mode 100644 index 0000000000..43de6da634 --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/passes/lower_associative_scan.py @@ -0,0 +1,258 @@ +"""Lower ``higher_order.associative_scan`` for Mamba's affine recurrence. + +``combine_mode="pointwise"`` keeps the HOP intact through export. The only +combine pattern this pass accepts is Mamba's first-order linear recurrence:: + + combine((a_l, b_l), (a_r, b_r)) = (a_l * a_r, a_r * b_l + b_r) + +which is ``h_t = a_t * h_{t-1} + b_t``. Anything else is left alone. + +For a static scan length the HOP is replaced with a Hillis–Steele inclusive +scan built from ``slice`` / ``cat`` / ``mul`` / ``add`` — all of which already +have converters. Dynamic sequence lengths are declined so today's PyTorch +fallback behaviour is preserved. +""" + +from __future__ import annotations + +import logging +import math +import operator +from typing import Optional, Tuple + +import torch +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( + clean_up_graph_after_modifications, +) + +logger = logging.getLogger(__name__) + + +def _is_associative_scan(target: object) -> bool: + if target is torch.ops.higher_order.associative_scan: + return True + name = str(target) + return "associative_scan" in name and "higher_order" in name + + +def _is_mul(n: torch.fx.Node, x: torch.fx.Node, y: torch.fx.Node) -> bool: + return ( + n.op == "call_function" + and n.target + in (torch.ops.aten.mul.Tensor, torch.ops.aten.mul.default, operator.mul) + and set(n.args[:2]) == {x, y} + ) + + +def _is_add(n: torch.fx.Node, x: torch.fx.Node, y: torch.fx.Node) -> bool: + return ( + n.op == "call_function" + and n.target + in (torch.ops.aten.add.Tensor, torch.ops.aten.add.default, operator.add) + and set(n.args[:2]) == {x, y} + ) + + +def _unwrap_output_pair( + output_node: torch.fx.Node, +) -> Optional[Tuple[torch.fx.Node, torch.fx.Node]]: + out_args = output_node.args[0] + # wrap_combine_fn_flat may nest the pair as a list/tuple of length 1 + while isinstance(out_args, (list, tuple)) and len(out_args) == 1: + out_args = out_args[0] + if not isinstance(out_args, (list, tuple)) or len(out_args) != 2: + return None + o0, o1 = out_args + if not isinstance(o0, torch.fx.Node) or not isinstance(o1, torch.fx.Node): + return None + return o0, o1 + + +def _is_mamba_affine_combine(combine_gm: torch.fx.GraphModule) -> bool: + """True iff the combine subgraph is ``(a_l*a_r, a_r*b_l + b_r)``.""" + placeholders = [n for n in combine_gm.graph.nodes if n.op == "placeholder"] + output = next(n for n in combine_gm.graph.nodes if n.op == "output") + if len(placeholders) != 4: + return False + + pair = _unwrap_output_pair(output) + if pair is None: + return False + o0, o1 = pair + + a_l, b_l, a_r, b_r = placeholders + if not _is_mul(o0, a_l, a_r): + return False + + # o1 = (a_r * b_l) + b_r + mul_nodes = [ + n + for n in combine_gm.graph.nodes + if n.op == "call_function" + and n.target + in (torch.ops.aten.mul.Tensor, torch.ops.aten.mul.default, operator.mul) + and set(n.args[:2]) == {a_r, b_l} + ] + if len(mul_nodes) != 1: + return False + return _is_add(o1, mul_nodes[0], b_r) + + +def _static_scan_length(xs: Tuple[torch.fx.Node, ...]) -> Optional[int]: + """Return the concrete length along dim 0, or None if dynamic/unknown.""" + lengths = set() + for x in xs: + val = x.meta.get("val") + if val is None or not hasattr(val, "shape") or len(val.shape) < 1: + return None + length = val.shape[0] + if isinstance(length, torch.SymInt): + return None + if not isinstance(length, int) or length < 0: + return None + lengths.add(int(length)) + if len(lengths) != 1: + return None + return lengths.pop() + + +def _hillis_steele_scan( + gm: torch.fx.GraphModule, + a: torch.fx.Node, + b: torch.fx.Node, + scan_len: int, + before: torch.fx.Node, +) -> Tuple[torch.fx.Node, torch.fx.Node]: + """Inclusive Hillis–Steele scan of ``(a, b)`` along dim 0. + + Each stage combines position ``i`` with ``i - step``. The leading ``step`` + positions have no left operand, so they are carried through unchanged. + """ + with gm.graph.inserting_before(before): + for d in range(math.ceil(math.log2(scan_len)) if scan_len > 1 else 0): + step = 1 << d + a_head = gm.graph.call_function( + torch.ops.aten.slice.Tensor, (a, 0, 0, step, 1) + ) + b_head = gm.graph.call_function( + torch.ops.aten.slice.Tensor, (b, 0, 0, step, 1) + ) + a_left = gm.graph.call_function( + torch.ops.aten.slice.Tensor, (a, 0, 0, scan_len - step, 1) + ) + b_left = gm.graph.call_function( + torch.ops.aten.slice.Tensor, (b, 0, 0, scan_len - step, 1) + ) + a_right = gm.graph.call_function( + torch.ops.aten.slice.Tensor, (a, 0, step, scan_len, 1) + ) + b_right = gm.graph.call_function( + torch.ops.aten.slice.Tensor, (b, 0, step, scan_len, 1) + ) + a_tail = gm.graph.call_function( + torch.ops.aten.mul.Tensor, (a_left, a_right) + ) + tmp = gm.graph.call_function(torch.ops.aten.mul.Tensor, (a_right, b_left)) + b_tail = gm.graph.call_function(torch.ops.aten.add.Tensor, (tmp, b_right)) + a = gm.graph.call_function( + torch.ops.aten.cat.default, ([a_head, a_tail], 0) + ) + b = gm.graph.call_function( + torch.ops.aten.cat.default, ([b_head, b_tail], 0) + ) + return a, b + + +def _getitem_index(user: torch.fx.Node) -> Optional[int]: + if user.op != "call_function": + return None + if user.target not in (operator.getitem, torch.ops.aten.select.int): + return None + idx = user.args[1] + if not isinstance(idx, int): + return None + return idx + + +def _rewrite_associative_scan(gm: torch.fx.GraphModule, node: torch.fx.Node) -> bool: + if len(node.args) < 2: + return False + + combine_node = node.args[0] + xs = node.args[1] + if combine_node.op != "get_attr": + return False + if not isinstance(xs, (list, tuple)) or len(xs) != 2: + return False + if any(not isinstance(x, torch.fx.Node) for x in xs): + return False + + # additional_inputs must be empty for the narrow Mamba pattern + if len(node.args) > 2 and node.args[2] not in ((), [], None): + return False + + try: + combine_gm = gm.get_submodule(combine_node.target) + except AttributeError: + combine_gm = getattr(gm, combine_node.target, None) + if not isinstance(combine_gm, torch.fx.GraphModule): + return False + if not _is_mamba_affine_combine(combine_gm): + logger.debug( + "associative_scan %s: combine subgraph is not the Mamba affine pattern", + node.name, + ) + return False + + scan_len = _static_scan_length(tuple(xs)) + if scan_len is None: + logger.debug( + "associative_scan %s: dynamic scan length; leaving HOP for PyTorch fallback", + node.name, + ) + return False + + # Validate users before mutating the graph. + replacements = {} + for user in list(node.users): + idx = _getitem_index(user) + if idx not in (0, 1): + logger.debug( + "associative_scan %s: unexpected user %s; leaving HOP alone", + node.name, + user, + ) + return False + replacements[user] = idx + + a_out, b_out = _hillis_steele_scan(gm, xs[0], xs[1], scan_len, node) + outs = (a_out, b_out) + for user, idx in replacements.items(): + user.replace_all_uses_with(outs[idx]) + gm.graph.erase_node(user) + gm.graph.erase_node(node) + logger.debug( + "associative_scan %s: lowered Mamba affine scan (S=%d) to Hillis-Steele aten ops", + node.name, + scan_len, + ) + return True + + +def lower_associative_scan( + gm: torch.fx.GraphModule, settings: CompilationSettings +) -> torch.fx.GraphModule: + """Replace Mamba-style ``associative_scan`` HOPs with a parallel aten scan.""" + changed = False + for node in list(gm.graph.nodes): + if node.op != "call_function" or not _is_associative_scan(node.target): + continue + if _rewrite_associative_scan(gm, node): + changed = True + + if changed: + gm = clean_up_graph_after_modifications(gm) + logger.debug("After lower_associative_scan:\n%s", gm.graph) + + return gm diff --git a/tests/py/dynamo/lowering/test_lower_associative_scan.py b/tests/py/dynamo/lowering/test_lower_associative_scan.py new file mode 100644 index 0000000000..c18c231c06 --- /dev/null +++ b/tests/py/dynamo/lowering/test_lower_associative_scan.py @@ -0,0 +1,215 @@ +# type: ignore +"""Tests for ``lower_associative_scan`` (Mamba affine recurrence). + +``combine_mode="pointwise"`` keeps ``higher_order.associative_scan`` through +export. The pass matches only the combine ``(a_l*a_r, a_r*b_l + b_r)`` with a +static scan length and replaces it with a Hillis–Steele aten scan. +""" + +import math +import unittest + +import torch +import torch_tensorrt +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, + pre_export_lowering, +) +from torch_tensorrt.dynamo.lowering.passes.lower_associative_scan import ( + _is_mamba_affine_combine, + lower_associative_scan, +) + + +def _mamba_scan_module(combine_mode: str) -> torch.nn.Module: + class MambaScan(torch.nn.Module): + def __init__(self, mode: str): + super().__init__() + self.combine_mode = mode + + def forward(self, discrete_a, delta_b_u, c): + from torch._higher_order_ops.associative_scan import associative_scan + + def combine_fn(left, right): + a_left, b_left = left + a_right, b_right = right + return (a_left * a_right, a_right * b_left + b_right) + + _, all_h = associative_scan( + combine_fn, + (discrete_a, delta_b_u), + dim=2, + combine_mode=self.combine_mode, + ) + return ( + torch.matmul(all_h.permute(0, 2, 1, 3), c.unsqueeze(-1)) + .squeeze(-1) + .permute(0, 2, 1) + ) + + return MambaScan(combine_mode) + + +def _has_associative_scan(gm: torch.fx.GraphModule) -> bool: + for n in gm.graph.nodes: + if n.op != "call_function": + continue + name = str(n.target) + if "associative_scan" in name and "higher_order" in name: + return True + if n.target is getattr(torch.ops.higher_order, "associative_scan", None): + return True + return False + + +def _pytorch_segments(compiled) -> list: + """Names of the segments the partitioner left in PyTorch.""" + return [name for name, _ in compiled.named_children() if "_run_on_gpu" in name] + + +def _lower_exported(model, inputs, experimental: bool = False): + settings = CompilationSettings(min_block_size=1) + with torch.no_grad(): + ep = torch.export.export(model, inputs) + ep = pre_export_lowering(ep, settings) + ep = ep.run_decompositions(get_decompositions(experimental)) + return post_lowering(ep.module(), settings), ep + + +@unittest.skipIf(not torch.cuda.is_available(), "CUDA required") +class TestLowerAssociativeScan(TestCase): + def test_pointwise_scan_removed_from_graph(self): + b, d, s, n = 1, 4, 8, 16 + model = _mamba_scan_module("pointwise").cuda().eval() + inputs = ( + torch.rand(b, d, s, n, device="cuda"), + torch.randn(b, d, s, n, device="cuda"), + torch.randn(b, s, n, device="cuda"), + ) + gm, _ = _lower_exported(model, inputs) + self.assertFalse( + _has_associative_scan(gm), + f"associative_scan still present after lowering:\n{gm.graph}", + ) + # Parallel scan stages should introduce doubling slices/cats. + targets = {n.target for n in gm.graph.nodes if n.op == "call_function"} + self.assertIn(torch.ops.aten.slice.Tensor, targets) + self.assertIn(torch.ops.aten.cat.default, targets) + self.assertIn(torch.ops.aten.mul.Tensor, targets) + self.assertIn(torch.ops.aten.add.Tensor, targets) + # ones_like / zeros_like have no converter after run_decompositions. + for op in (torch.ops.aten.ones_like.default, torch.ops.aten.zeros_like.default): + self.assertNotIn(op, targets) + + def test_pointwise_scan_numerics_match_eager(self): + b, d, s, n = 1, 2, 8, 4 + model = _mamba_scan_module("pointwise").cuda().eval() + inputs = ( + torch.rand(b, d, s, n, device="cuda"), + torch.randn(b, d, s, n, device="cuda"), + torch.randn(b, s, n, device="cuda"), + ) + gm, ep = _lower_exported(model, inputs) + self.assertFalse(_has_associative_scan(gm)) + + ref = model(*[t.clone() for t in inputs]) + # Execute the lowered GraphModule (aten scan, no HOP). + out = gm(*[t.clone() for t in inputs]) + torch.testing.assert_close(out, ref, rtol=1e-4, atol=1e-4) + + compiled = torch_tensorrt.dynamo.compile( + ep, + inputs=list(inputs), + enabled_precisions={torch.float32}, + min_block_size=1, + ) + self.assertEqual( + _pytorch_segments(compiled), [], "the scan must run entirely in TRT" + ) + trt_out = compiled(*[t.clone() for t in inputs]) + torch.testing.assert_close(trt_out, ref, rtol=1e-3, atol=1e-3) + + def test_non_power_of_two_length(self): + b, d, s, n = 1, 2, 7, 4 + model = _mamba_scan_module("pointwise").cuda().eval() + inputs = ( + torch.rand(b, d, s, n, device="cuda"), + torch.randn(b, d, s, n, device="cuda"), + torch.randn(b, s, n, device="cuda"), + ) + gm, ep = _lower_exported(model, inputs) + self.assertFalse(_has_associative_scan(gm)) + + ref = model(*[t.clone() for t in inputs]) + torch.testing.assert_close( + gm(*[t.clone() for t in inputs]), ref, rtol=1e-4, atol=1e-4 + ) + # One add per stage; a-chain cats are DCE'd since only b is consumed. + adds = [n for n in gm.graph.nodes if n.target is torch.ops.aten.add.Tensor] + self.assertEqual(len(adds), math.ceil(math.log2(s))) + + compiled = torch_tensorrt.dynamo.compile( + ep, + inputs=list(inputs), + enabled_precisions={torch.float32}, + min_block_size=1, + ) + self.assertEqual( + _pytorch_segments(compiled), [], "the scan must run entirely in TRT" + ) + torch.testing.assert_close( + compiled(*[t.clone() for t in inputs]), ref, rtol=1e-3, atol=1e-3 + ) + + def test_decline_non_mamba_combine(self): + """A different associative combine must keep the HOP.""" + + class SumScan(torch.nn.Module): + def forward(self, x): + from torch._higher_order_ops.associative_scan import associative_scan + + def combine_fn(left, right): + return left + right + + return associative_scan( + combine_fn, x, dim=0, combine_mode="pointwise" + ) + + model = SumScan().cuda().eval() + x = torch.randn(8, 4, device="cuda") + settings = CompilationSettings(min_block_size=1) + with torch.no_grad(): + ep = torch.export.export(model, (x,)) + ep = pre_export_lowering(ep, settings) + ep = ep.run_decompositions(get_decompositions(False)) + gm = ep.module() + + # Confirm a scan HOP is present, then that our pass declines it. + self.assertTrue(_has_associative_scan(gm)) + gm2 = lower_associative_scan(gm, settings) + self.assertTrue(_has_associative_scan(gm2)) + + def test_mamba_combine_matcher_unit(self): + """Direct unit check of the narrow combine matcher.""" + + class Combine(torch.nn.Module): + def forward(self, a_l, b_l, a_r, b_r): + return a_l * a_r, a_r * b_l + b_r + + traced = torch.fx.symbolic_trace(Combine()) + self.assertTrue(_is_mamba_affine_combine(traced)) + + class BadCombine(torch.nn.Module): + def forward(self, a_l, b_l, a_r, b_r): + return a_l + a_r, b_l + b_r + + self.assertFalse( + _is_mamba_affine_combine(torch.fx.symbolic_trace(BadCombine())) + ) + + +if __name__ == "__main__": + run_tests()