diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index c4f9019e59..c0c2481d86 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -3665,6 +3665,101 @@ def aten_ops_sort( ) +def _arg_or_kwarg( + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + index: int, + key: str, + default: Any = None, +) -> Any: + """Read an operand that export may have left positional or keyword.""" + if key in kwargs and kwargs[key] is not None: + return kwargs[key] + return args_bounds_check(args, index, default) + + +def histc_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + bins = _arg_or_kwarg(node.args, node.kwargs, 1, "bins", 100) + lo = _arg_or_kwarg(node.args, node.kwargs, 2, "min", 0) + hi = _arg_or_kwarg(node.args, node.kwargs, 3, "max", 0) + if not all( + isinstance(v, (int, float)) and not isinstance(v, bool) for v in (bins, lo, hi) + ): + return False + # min == max == 0 asks aten to take the bin edges from the data itself. + return int(bins) > 0 and float(lo) < float(hi) + + +@dynamo_tensorrt_converter( + torch.ops.aten.histc.default, capability_validator=histc_validator +) +@enforce_tensor_types({0: (TRTTensor,)}) +def aten_ops_histc( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + return impl.moe.histc( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + int(_arg_or_kwarg(args, kwargs, 1, "bins", 100)), + float(_arg_or_kwarg(args, kwargs, 2, "min", 0)), + float(_arg_or_kwarg(args, kwargs, 3, "max", 0)), + ) + + +def grouped_mm_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + # Per-group bias would need a gather this lowering does not implement. + if _arg_or_kwarg(node.args, node.kwargs, 3, "bias") is not None: + return False + # Without offsets this is a plain bmm, which already has a converter. + if _arg_or_kwarg(node.args, node.kwargs, 2, "offs") is None: + return False + shapes = [] + for arg in node.args[:2]: + val = arg.meta.get("val") + if val is None: + return False + try: + shapes.append(tuple(int(d) for d in val.shape)) + except (TypeError, ValueError): # SymInt + return False + # Only the (rows, K) x (experts, K, N) form. + return len(shapes[0]) == 2 and len(shapes[1]) == 3 + + +@dynamo_tensorrt_converter( + torch.ops.aten._grouped_mm.default, capability_validator=grouped_mm_validator +) +@enforce_tensor_types({0: (TRTTensor,), 1: (TRTTensor,)}) +def aten_ops_grouped_mm( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + return impl.moe.grouped_mm( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + args[1], + _arg_or_kwarg(args, kwargs, 2, "offs"), + _arg_or_kwarg(args, kwargs, 4, "out_dtype"), + ) + + @dynamo_tensorrt_converter(torch.ops.aten.trunc.default, supports_dynamic_shapes=True) @enforce_tensor_types( { diff --git a/py/torch_tensorrt/dynamo/conversion/impl/__init__.py b/py/torch_tensorrt/dynamo/conversion/impl/__init__.py index 0cc92c40d4..d0e26cb994 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/__init__.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/__init__.py @@ -16,6 +16,7 @@ index_copy, linear, matmul, + moe, nccl_ops, normalization, pad, diff --git a/py/torch_tensorrt/dynamo/conversion/impl/moe.py b/py/torch_tensorrt/dynamo/conversion/impl/moe.py new file mode 100644 index 0000000000..531bf591c1 --- /dev/null +++ b/py/torch_tensorrt/dynamo/conversion/impl/moe.py @@ -0,0 +1,153 @@ +from typing import Optional, Union + +import numpy as np +import tensorrt as trt +import torch +from torch.fx.node import Target +from torch_tensorrt import _enums +from torch_tensorrt.dynamo._SourceIR import SourceIR +from torch_tensorrt.dynamo.conversion import impl +from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext +from torch_tensorrt.dynamo.conversion.converter_utils import ( + cast_trt_tensor, + get_trt_tensor, +) + + +def histc( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + input: trt.ITensor, + bins: int, + min: float, + max: float, +) -> trt.ITensor: + # aten returns the input dtype, so integer ids in means integer counts out. + out_dtype = _enums.dtype._from(input.dtype).to(torch.dtype) + + flat = impl.shuffle.reshape(ctx, target, source_ir, f"{name}_flat", input, (-1, 1)) + x = cast_trt_tensor(ctx, flat, torch.float32, f"{name}_f32", target, source_ir) + + # bin = floor((x - min) / (max - min) * bins), with x == max folded back into + # the last bin instead of overflowing to index `bins`. + shifted = impl.elementwise.sub(ctx, target, source_ir, f"{name}_shift", x, min) + scaled = impl.elementwise.mul( + ctx, target, source_ir, f"{name}_scale", shifted, bins / (max - min) + ) + idx = impl.unary.floor(ctx, target, source_ir, f"{name}_floor", scaled) + idx = impl.elementwise.min( + ctx, target, source_ir, f"{name}_lastbin", idx, float(bins - 1) + ) + + # Values outside [min, max] are dropped, not clamped. + in_range = impl.elementwise.logical_and( + ctx, + target, + source_ir, + f"{name}_inrange", + impl.elementwise.ge(ctx, target, source_ir, f"{name}_ge_min", x, min), + impl.elementwise.le(ctx, target, source_ir, f"{name}_le_max", x, max), + ) + + edges = get_trt_tensor( + ctx, np.arange(bins, dtype=np.float32).reshape(1, bins), f"{name}_edges" + ) + hit = impl.elementwise.eq(ctx, target, source_ir, f"{name}_hit", idx, edges) + hit = impl.elementwise.logical_and( + ctx, target, source_ir, f"{name}_hit_valid", hit, in_range + ) + + counts = impl.reduce.sum( + ctx, + target, + source_ir, + f"{name}_count", + cast_trt_tensor(ctx, hit, torch.float32, f"{name}_hit_f32", target, source_ir), + dim=0, + keepdim=False, + ) + return cast_trt_tensor(ctx, counts, out_dtype, f"{name}_out", target, source_ir) + + +def grouped_mm( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + mat1: trt.ITensor, + mat2: trt.ITensor, + offs: trt.ITensor, + out_dtype: Optional[Union[torch.dtype, trt.DataType]] = None, +) -> trt.ITensor: + num_rows = int(mat1.shape[0]) + num_experts = int(mat2.shape[0]) + compute_dtype = _enums.dtype._from(mat2.dtype).to(torch.dtype) + + # Rows arrive sorted by group and `offs` holds each group's exclusive end, + # so the group owning row i is how many offsets are <= i. + rows = get_trt_tensor( + ctx, np.arange(num_rows, dtype=np.int32).reshape(num_rows, 1), f"{name}_rows" + ) + offs_i32 = cast_trt_tensor( + ctx, offs, torch.int32, f"{name}_offs_i32", target, source_ir + ) + offs_row = impl.shuffle.reshape( + ctx, target, source_ir, f"{name}_offs_row", offs_i32, (1, -1) + ) + passed = impl.elementwise.le(ctx, target, source_ir, f"{name}_passed", offs_row, rows) + group = impl.reduce.sum( + ctx, + target, + source_ir, + f"{name}_group", + cast_trt_tensor( + ctx, passed, torch.int32, f"{name}_passed_i32", target, source_ir + ), + dim=1, + keepdim=False, + ) + + # Dense lowering: every row through every expert, then mask. + lhs = impl.shuffle.reshape( + ctx, target, source_ir, f"{name}_lhs", mat1, (1, num_rows, -1) + ) + per_expert = impl.matmul.matrix_multiply( + ctx, target, source_ir, f"{name}_mm", lhs, mat2 + ) + + expert_ids = get_trt_tensor( + ctx, + np.arange(num_experts, dtype=np.int32).reshape(num_experts, 1), + f"{name}_expert_ids", + ) + group_row = impl.shuffle.reshape( + ctx, target, source_ir, f"{name}_group_row", group, (1, num_rows) + ) + selected = impl.elementwise.eq( + ctx, target, source_ir, f"{name}_selected", expert_ids, group_row + ) + selected = impl.shuffle.reshape( + ctx, + target, + source_ir, + f"{name}_selected_3d", + cast_trt_tensor( + ctx, selected, compute_dtype, f"{name}_selected_cast", target, source_ir + ), + (num_experts, num_rows, 1), + ) + + masked = impl.elementwise.mul( + ctx, target, source_ir, f"{name}_masked", per_expert, selected + ) + out = impl.reduce.sum( + ctx, target, source_ir, f"{name}_reduce", masked, dim=0, keepdim=False + ) + + if out_dtype is not None: + out = cast_trt_tensor( + ctx, out, out_dtype, f"{name}_out_dtype", target, source_ir + ) + return out diff --git a/tests/py/dynamo/conversion/test_moe_aten.py b/tests/py/dynamo/conversion/test_moe_aten.py new file mode 100644 index 0000000000..e923496e47 --- /dev/null +++ b/tests/py/dynamo/conversion/test_moe_aten.py @@ -0,0 +1,74 @@ +import torch +import torch.nn as nn +from parameterized import parameterized +from torch.testing._internal.common_utils import run_tests + +from .harness import DispatchTestCase + + +class TestHistcConverter(DispatchTestCase): + @parameterized.expand( + [ + ((32,), 8, 0.0, 7.0), + ((16, 4), 4, -1.0, 2.0), + ((64,), 16, 0.0, 15.0), + ] + ) + def test_histc(self, shape, bins, lo, hi): + class Histc(nn.Module): + def forward(self, x): + return torch.ops.aten.histc.default(x, bins, lo, hi) + + inputs = [torch.rand(*shape) * (hi - lo + 2) + lo - 1] + self.run_test(Histc(), inputs, use_dynamo_tracer=True) + + def test_histc_integer_ids(self): + class Histc(nn.Module): + def forward(self, x): + return torch.ops.aten.histc.default(x, 8, 0, 7) + + inputs = [torch.randint(0, 8, (64,), dtype=torch.int32)] + self.run_test(Histc(), inputs, use_dynamo_tracer=True) + + +class TestGroupedMMConverter(DispatchTestCase): + @parameterized.expand( + [ + (16, 32, 24, 4, torch.bfloat16), + (32, 64, 48, 8, torch.bfloat16), + ] + ) + def test_grouped_mm(self, rows, k, n, experts, dtype): + class GroupedMM(nn.Module): + def forward(self, x, w, offs): + return torch.ops.aten._grouped_mm.default(x, w, offs) + + counts = torch.randint(0, max(rows // experts, 1) + 1, (experts,)) + offs = torch.cumsum(counts, 0).clamp(max=rows).to(torch.int32) + x = torch.randn(rows, k, dtype=dtype) + w = torch.randn(experts, k, n, dtype=dtype) + self.run_test( + GroupedMM(), + [x, w, offs], + use_dynamo_tracer=True, + ) + + def test_grouped_mm_kwargs_offs(self): + class GroupedMM(nn.Module): + def forward(self, x, w, offs): + return torch.ops.aten._grouped_mm.default(x, w, offs=offs) + + rows, k, n, experts = 16, 32, 24, 4 + counts = torch.tensor([3, 5, 0, 4]) + offs = torch.cumsum(counts, 0).clamp(max=rows).to(torch.int32) + x = torch.randn(rows, k, dtype=torch.bfloat16) + w = torch.randn(experts, k, n, dtype=torch.bfloat16) + self.run_test( + GroupedMM(), + [x, w, offs], + use_dynamo_tracer=True, + ) + + +if __name__ == "__main__": + run_tests()