MoE histc / _grouped_mm

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(
     {

py/torch_tensorrt/dynamo/conversion/impl/__init__.py

@@ -16,6 +16,7 @@
     index_copy,
     linear,
     matmul,
+    moe,
     nccl_ops,
     normalization,
     pad,

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

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()