fix: bump builds

build.toml

@@ -21,7 +21,9 @@ cuda-capabilities = [
     "9.0",
     "10.0",
     "10.1",
-    "12.0",
+    "11.8",
+    "12.0"
+    # "12.4"
 ]
 depends = ["torch", "cutlass_3_8"]
 src = [

build/torch26-cxx11-cu118-x86_64-linux/megablocks/__init__.py

@@ -5,11 +5,15 @@ import torch
 
 from ._ops import ops
 
-from megablocks.layers.arguments import Arguments
-from megablocks.layers.dmoe import ParallelDroplessMLP, dMoE
-from megablocks.layers.glu import SparseGLU
-from megablocks.layers.mlp import MLP, SparseMLP
-from megablocks.layers.moe import MoE, ParallelMLP, get_load_balancing_loss
+from .grouped_gemm import backend as gg_backend
+from .grouped_gemm import ops as gg_ops
+
+
+from .layers.arguments import Arguments
+from .layers.dmoe import ParallelDroplessMLP, dMoE
+from .layers.glu import SparseGLU
+from .layers.mlp import MLP, SparseMLP
+from .layers.moe import MoE, ParallelMLP, get_load_balancing_loss
 
 # This section contains the direct kernel exports (not inlcuded in the original code)
 def exclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor:

build/torch26-cxx11-cu118-x86_64-linux/megablocks/{_megablocks_359242d.abi3.so → _megablocks_a585153_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:36a918d61308a0acbc880516a42fcc2c4dc393f25655326e52128465c9501709
-size 10456376
+oid sha256:44462d45f75616c369c2421fe41d53cd1d1dc365f1d2545d870e2db999e67e38
+size 10517608

build/torch26-cxx11-cu118-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_359242d
-ops = torch.ops._megablocks_359242d
+from . import _megablocks_a585153_dirty
+ops = torch.ops._megablocks_a585153_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_359242d::{op_name}"
+    return f"_megablocks_a585153_dirty::{op_name}"

build/torch26-cxx11-cu118-x86_64-linux/megablocks/grouped_gemm/__init__.py

@@ -0,0 +1,2 @@
+from . import ops
+from . import backend

build/torch26-cxx11-cu118-x86_64-linux/megablocks/grouped_gemm/backend.py

@@ -0,0 +1,32 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+
+# # TODO(tgale): Wrap this in a try-block with better
+# # error message and instructions for building the
+# # c++ operations.
+# import grouped_gemm_backend as backend
+
+# We import the backend operations from the megablocks package as
+# grouped_gemm is vendored in megablocks in this repository.
+# from ... import _ops as backend
+from megablocks._ops import ops as backend  # type: ignore
+
+def _allocate_output(a, b, batch_sizes, trans_a, trans_b):
+    assert not (trans_a and trans_b)
+    assert batch_sizes.ndim == 1, "Expected 1d tensor for batch_sizes"
+    assert a.ndim == 2, "Expected 2d tensor for 'a'"
+    assert b.ndim == (2 if trans_a else 3)
+
+    shape = (
+        (batch_sizes.shape[0], a.shape[1], b.shape[1])
+        if trans_a else
+        (a.shape[0], (b.shape[1] if trans_b else b.shape[2]))
+    )
+    return torch.empty(*shape, device=a.device, dtype=a.dtype)
+
+def gmm(a, b, batch_sizes, trans_a=False, trans_b=False, c=None):
+    if c is None:
+        c = _allocate_output(a, b, batch_sizes, trans_a, trans_b)
+    backend.gmm(a, b, c, batch_sizes, trans_a, trans_b)
+    return c

build/torch26-cxx11-cu118-x86_64-linux/megablocks/grouped_gemm/ops.py

@@ -0,0 +1,33 @@
+from . import backend
+import torch
+
+
+class GroupedGemm(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, a, b, batch_sizes, trans_b):
+        ctx.save_for_backward(a, b, batch_sizes)
+        ctx.trans_b = trans_b
+        return backend.gmm(a, b, batch_sizes, trans_a=False, trans_b=trans_b)
+
+    @staticmethod
+    def backward(ctx, grad):
+        grad = grad.contiguous()
+        a, b, batch_sizes = ctx.saved_tensors
+        trans_b = ctx.trans_b
+
+        agrad = None
+        if ctx.needs_input_grad[0]:
+            agrad = backend.gmm(
+                grad, b, batch_sizes, trans_a=False, trans_b=not trans_b)
+
+        bgrad = None
+        if ctx.needs_input_grad[1]:
+            lhs, rhs = (grad, a) if trans_b else (a, grad)
+            bgrad = backend.gmm(
+                lhs, rhs, batch_sizes, trans_a=True, trans_b=False)
+        return agrad, bgrad, None, None
+
+
+def gmm(a, b, batch_sizes, trans_b=False):
+    return GroupedGemm.apply(a, b, batch_sizes, trans_b)

build/torch26-cxx11-cu118-x86_64-linux/megablocks/grouped_gemm_util.py

@@ -4,7 +4,8 @@ import warnings
 
 _grouped_gemm_is_available: bool = False
 try:
-    import grouped_gemm
+    # import grouped_gemm
+    pass
     _grouped_gemm_is_available = True
 except ImportError as error:
     warnings.warn('Grouped GEMM not available.')
@@ -22,5 +23,9 @@ def assert_grouped_gemm_is_available():
     assert _grouped_gemm_is_available, msg
 
 
-backend = grouped_gemm.backend if grouped_gemm_is_available() else None
-ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+# backend = grouped_gemm.backend if grouped_gemm_is_available() else None
+# ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+
+
+from .grouped_gemm import backend as ops
+from .grouped_gemm import ops as backend

build/torch26-cxx11-cu118-x86_64-linux/megablocks/layers/__init__.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # from megablocks.layers.dmoe import dMoE
-from megablocks.layers.moe import MoE
+from .moe import MoE
 
 __all__ = [
     'MoE',

build/torch26-cxx11-cu124-x86_64-linux/megablocks/__init__.py

@@ -5,11 +5,15 @@ import torch
 
 from ._ops import ops
 
-from megablocks.layers.arguments import Arguments
-from megablocks.layers.dmoe import ParallelDroplessMLP, dMoE
-from megablocks.layers.glu import SparseGLU
-from megablocks.layers.mlp import MLP, SparseMLP
-from megablocks.layers.moe import MoE, ParallelMLP, get_load_balancing_loss
+from .grouped_gemm import backend as gg_backend
+from .grouped_gemm import ops as gg_ops
+
+
+from .layers.arguments import Arguments
+from .layers.dmoe import ParallelDroplessMLP, dMoE
+from .layers.glu import SparseGLU
+from .layers.mlp import MLP, SparseMLP
+from .layers.moe import MoE, ParallelMLP, get_load_balancing_loss
 
 # This section contains the direct kernel exports (not inlcuded in the original code)
 def exclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor:

build/torch26-cxx11-cu124-x86_64-linux/megablocks/{_megablocks_359242d.abi3.so → _megablocks_a585153_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:12cf047c4bcb5f368f490ada3dafcecd1242a443ff5ded94c09d62548922098c
-size 11795992
+oid sha256:e734576700345e035790357ea19730e84e90c176747076ce845995bc3a0e0d50
+size 11869424

build/torch26-cxx11-cu124-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_359242d
-ops = torch.ops._megablocks_359242d
+from . import _megablocks_a585153_dirty
+ops = torch.ops._megablocks_a585153_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_359242d::{op_name}"
+    return f"_megablocks_a585153_dirty::{op_name}"

build/torch26-cxx11-cu124-x86_64-linux/megablocks/grouped_gemm/__init__.py

@@ -0,0 +1,2 @@
+from . import ops
+from . import backend

build/torch26-cxx11-cu124-x86_64-linux/megablocks/grouped_gemm/backend.py

@@ -0,0 +1,32 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+
+# # TODO(tgale): Wrap this in a try-block with better
+# # error message and instructions for building the
+# # c++ operations.
+# import grouped_gemm_backend as backend
+
+# We import the backend operations from the megablocks package as
+# grouped_gemm is vendored in megablocks in this repository.
+# from ... import _ops as backend
+from megablocks._ops import ops as backend  # type: ignore
+
+def _allocate_output(a, b, batch_sizes, trans_a, trans_b):
+    assert not (trans_a and trans_b)
+    assert batch_sizes.ndim == 1, "Expected 1d tensor for batch_sizes"
+    assert a.ndim == 2, "Expected 2d tensor for 'a'"
+    assert b.ndim == (2 if trans_a else 3)
+
+    shape = (
+        (batch_sizes.shape[0], a.shape[1], b.shape[1])
+        if trans_a else
+        (a.shape[0], (b.shape[1] if trans_b else b.shape[2]))
+    )
+    return torch.empty(*shape, device=a.device, dtype=a.dtype)
+
+def gmm(a, b, batch_sizes, trans_a=False, trans_b=False, c=None):
+    if c is None:
+        c = _allocate_output(a, b, batch_sizes, trans_a, trans_b)
+    backend.gmm(a, b, c, batch_sizes, trans_a, trans_b)
+    return c

build/torch26-cxx11-cu124-x86_64-linux/megablocks/grouped_gemm/ops.py

@@ -0,0 +1,33 @@
+from . import backend
+import torch
+
+
+class GroupedGemm(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, a, b, batch_sizes, trans_b):
+        ctx.save_for_backward(a, b, batch_sizes)
+        ctx.trans_b = trans_b
+        return backend.gmm(a, b, batch_sizes, trans_a=False, trans_b=trans_b)
+
+    @staticmethod
+    def backward(ctx, grad):
+        grad = grad.contiguous()
+        a, b, batch_sizes = ctx.saved_tensors
+        trans_b = ctx.trans_b
+
+        agrad = None
+        if ctx.needs_input_grad[0]:
+            agrad = backend.gmm(
+                grad, b, batch_sizes, trans_a=False, trans_b=not trans_b)
+
+        bgrad = None
+        if ctx.needs_input_grad[1]:
+            lhs, rhs = (grad, a) if trans_b else (a, grad)
+            bgrad = backend.gmm(
+                lhs, rhs, batch_sizes, trans_a=True, trans_b=False)
+        return agrad, bgrad, None, None
+
+
+def gmm(a, b, batch_sizes, trans_b=False):
+    return GroupedGemm.apply(a, b, batch_sizes, trans_b)

build/torch26-cxx11-cu124-x86_64-linux/megablocks/grouped_gemm_util.py

@@ -4,7 +4,8 @@ import warnings
 
 _grouped_gemm_is_available: bool = False
 try:
-    import grouped_gemm
+    # import grouped_gemm
+    pass
     _grouped_gemm_is_available = True
 except ImportError as error:
     warnings.warn('Grouped GEMM not available.')
@@ -22,5 +23,9 @@ def assert_grouped_gemm_is_available():
     assert _grouped_gemm_is_available, msg
 
 
-backend = grouped_gemm.backend if grouped_gemm_is_available() else None
-ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+# backend = grouped_gemm.backend if grouped_gemm_is_available() else None
+# ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+
+
+from .grouped_gemm import backend as ops
+from .grouped_gemm import ops as backend

build/torch26-cxx11-cu124-x86_64-linux/megablocks/layers/__init__.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # from megablocks.layers.dmoe import dMoE
-from megablocks.layers.moe import MoE
+from .moe import MoE
 
 __all__ = [
     'MoE',

build/torch26-cxx11-cu126-x86_64-linux/megablocks/__init__.py

@@ -5,11 +5,15 @@ import torch
 
 from ._ops import ops
 
-from megablocks.layers.arguments import Arguments
-from megablocks.layers.dmoe import ParallelDroplessMLP, dMoE
-from megablocks.layers.glu import SparseGLU
-from megablocks.layers.mlp import MLP, SparseMLP
-from megablocks.layers.moe import MoE, ParallelMLP, get_load_balancing_loss
+from .grouped_gemm import backend as gg_backend
+from .grouped_gemm import ops as gg_ops
+
+
+from .layers.arguments import Arguments
+from .layers.dmoe import ParallelDroplessMLP, dMoE
+from .layers.glu import SparseGLU
+from .layers.mlp import MLP, SparseMLP
+from .layers.moe import MoE, ParallelMLP, get_load_balancing_loss
 
 # This section contains the direct kernel exports (not inlcuded in the original code)
 def exclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor:

build/torch26-cxx11-cu126-x86_64-linux/megablocks/{_megablocks_359242d.abi3.so → _megablocks_a585153_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:38fcc4266dd94ee3f307bebd1264a1faaf91c73dc680adb72cd268748130b10f
-size 11835888
+oid sha256:8507dd1e6fc8f4df45af233d506ef96b962cacecf9e2d0694247547b0dd7dde0
+size 11931080

build/torch26-cxx11-cu126-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_359242d
-ops = torch.ops._megablocks_359242d
+from . import _megablocks_a585153_dirty
+ops = torch.ops._megablocks_a585153_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_359242d::{op_name}"
+    return f"_megablocks_a585153_dirty::{op_name}"

build/torch26-cxx11-cu126-x86_64-linux/megablocks/grouped_gemm/__init__.py

@@ -0,0 +1,2 @@
+from . import ops
+from . import backend

build/torch26-cxx11-cu126-x86_64-linux/megablocks/grouped_gemm/backend.py

@@ -0,0 +1,32 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+
+# # TODO(tgale): Wrap this in a try-block with better
+# # error message and instructions for building the
+# # c++ operations.
+# import grouped_gemm_backend as backend
+
+# We import the backend operations from the megablocks package as
+# grouped_gemm is vendored in megablocks in this repository.
+# from ... import _ops as backend
+from megablocks._ops import ops as backend  # type: ignore
+
+def _allocate_output(a, b, batch_sizes, trans_a, trans_b):
+    assert not (trans_a and trans_b)
+    assert batch_sizes.ndim == 1, "Expected 1d tensor for batch_sizes"
+    assert a.ndim == 2, "Expected 2d tensor for 'a'"
+    assert b.ndim == (2 if trans_a else 3)
+
+    shape = (
+        (batch_sizes.shape[0], a.shape[1], b.shape[1])
+        if trans_a else
+        (a.shape[0], (b.shape[1] if trans_b else b.shape[2]))
+    )
+    return torch.empty(*shape, device=a.device, dtype=a.dtype)
+
+def gmm(a, b, batch_sizes, trans_a=False, trans_b=False, c=None):
+    if c is None:
+        c = _allocate_output(a, b, batch_sizes, trans_a, trans_b)
+    backend.gmm(a, b, c, batch_sizes, trans_a, trans_b)
+    return c

build/torch26-cxx11-cu126-x86_64-linux/megablocks/grouped_gemm/ops.py

@@ -0,0 +1,33 @@
+from . import backend
+import torch
+
+
+class GroupedGemm(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, a, b, batch_sizes, trans_b):
+        ctx.save_for_backward(a, b, batch_sizes)
+        ctx.trans_b = trans_b
+        return backend.gmm(a, b, batch_sizes, trans_a=False, trans_b=trans_b)
+
+    @staticmethod
+    def backward(ctx, grad):
+        grad = grad.contiguous()
+        a, b, batch_sizes = ctx.saved_tensors
+        trans_b = ctx.trans_b
+
+        agrad = None
+        if ctx.needs_input_grad[0]:
+            agrad = backend.gmm(
+                grad, b, batch_sizes, trans_a=False, trans_b=not trans_b)
+
+        bgrad = None
+        if ctx.needs_input_grad[1]:
+            lhs, rhs = (grad, a) if trans_b else (a, grad)
+            bgrad = backend.gmm(
+                lhs, rhs, batch_sizes, trans_a=True, trans_b=False)
+        return agrad, bgrad, None, None
+
+
+def gmm(a, b, batch_sizes, trans_b=False):
+    return GroupedGemm.apply(a, b, batch_sizes, trans_b)

build/torch26-cxx11-cu126-x86_64-linux/megablocks/grouped_gemm_util.py

@@ -4,7 +4,8 @@ import warnings
 
 _grouped_gemm_is_available: bool = False
 try:
-    import grouped_gemm
+    # import grouped_gemm
+    pass
     _grouped_gemm_is_available = True
 except ImportError as error:
     warnings.warn('Grouped GEMM not available.')
@@ -22,5 +23,9 @@ def assert_grouped_gemm_is_available():
     assert _grouped_gemm_is_available, msg
 
 
-backend = grouped_gemm.backend if grouped_gemm_is_available() else None
-ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+# backend = grouped_gemm.backend if grouped_gemm_is_available() else None
+# ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+
+
+from .grouped_gemm import backend as ops
+from .grouped_gemm import ops as backend

build/torch26-cxx11-cu126-x86_64-linux/megablocks/layers/__init__.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # from megablocks.layers.dmoe import dMoE
-from megablocks.layers.moe import MoE
+from .moe import MoE
 
 __all__ = [
     'MoE',

build/torch26-cxx98-cu118-x86_64-linux/megablocks/__init__.py

@@ -5,11 +5,15 @@ import torch
 
 from ._ops import ops
 
-from megablocks.layers.arguments import Arguments
-from megablocks.layers.dmoe import ParallelDroplessMLP, dMoE
-from megablocks.layers.glu import SparseGLU
-from megablocks.layers.mlp import MLP, SparseMLP
-from megablocks.layers.moe import MoE, ParallelMLP, get_load_balancing_loss
+from .grouped_gemm import backend as gg_backend
+from .grouped_gemm import ops as gg_ops
+
+
+from .layers.arguments import Arguments
+from .layers.dmoe import ParallelDroplessMLP, dMoE
+from .layers.glu import SparseGLU
+from .layers.mlp import MLP, SparseMLP
+from .layers.moe import MoE, ParallelMLP, get_load_balancing_loss
 
 # This section contains the direct kernel exports (not inlcuded in the original code)
 def exclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor:

build/torch26-cxx98-cu118-x86_64-linux/megablocks/{_megablocks_359242d.abi3.so → _megablocks_a585153_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c9661a5704bd53f129788d2b9241f44cf00e1447dab8127d5a07675b1d6ca2ba
-size 10444224
+oid sha256:6dc0dcea20fc1350689addf7cb9927f7bb709f68ed89d4c711b0f7db579a463b
+size 10510072

build/torch26-cxx98-cu118-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_359242d
-ops = torch.ops._megablocks_359242d
+from . import _megablocks_a585153_dirty
+ops = torch.ops._megablocks_a585153_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_359242d::{op_name}"
+    return f"_megablocks_a585153_dirty::{op_name}"

build/torch26-cxx98-cu118-x86_64-linux/megablocks/grouped_gemm/__init__.py

@@ -0,0 +1,2 @@
+from . import ops
+from . import backend

build/torch26-cxx98-cu118-x86_64-linux/megablocks/grouped_gemm/backend.py

@@ -0,0 +1,32 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+
+# # TODO(tgale): Wrap this in a try-block with better
+# # error message and instructions for building the
+# # c++ operations.
+# import grouped_gemm_backend as backend
+
+# We import the backend operations from the megablocks package as
+# grouped_gemm is vendored in megablocks in this repository.
+# from ... import _ops as backend
+from megablocks._ops import ops as backend  # type: ignore
+
+def _allocate_output(a, b, batch_sizes, trans_a, trans_b):
+    assert not (trans_a and trans_b)
+    assert batch_sizes.ndim == 1, "Expected 1d tensor for batch_sizes"
+    assert a.ndim == 2, "Expected 2d tensor for 'a'"
+    assert b.ndim == (2 if trans_a else 3)
+
+    shape = (
+        (batch_sizes.shape[0], a.shape[1], b.shape[1])
+        if trans_a else
+        (a.shape[0], (b.shape[1] if trans_b else b.shape[2]))
+    )
+    return torch.empty(*shape, device=a.device, dtype=a.dtype)
+
+def gmm(a, b, batch_sizes, trans_a=False, trans_b=False, c=None):
+    if c is None:
+        c = _allocate_output(a, b, batch_sizes, trans_a, trans_b)
+    backend.gmm(a, b, c, batch_sizes, trans_a, trans_b)
+    return c

build/torch26-cxx98-cu118-x86_64-linux/megablocks/grouped_gemm/ops.py

@@ -0,0 +1,33 @@
+from . import backend
+import torch
+
+
+class GroupedGemm(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, a, b, batch_sizes, trans_b):
+        ctx.save_for_backward(a, b, batch_sizes)
+        ctx.trans_b = trans_b
+        return backend.gmm(a, b, batch_sizes, trans_a=False, trans_b=trans_b)
+
+    @staticmethod
+    def backward(ctx, grad):
+        grad = grad.contiguous()
+        a, b, batch_sizes = ctx.saved_tensors
+        trans_b = ctx.trans_b
+
+        agrad = None
+        if ctx.needs_input_grad[0]:
+            agrad = backend.gmm(
+                grad, b, batch_sizes, trans_a=False, trans_b=not trans_b)
+
+        bgrad = None
+        if ctx.needs_input_grad[1]:
+            lhs, rhs = (grad, a) if trans_b else (a, grad)
+            bgrad = backend.gmm(
+                lhs, rhs, batch_sizes, trans_a=True, trans_b=False)
+        return agrad, bgrad, None, None
+
+
+def gmm(a, b, batch_sizes, trans_b=False):
+    return GroupedGemm.apply(a, b, batch_sizes, trans_b)

build/torch26-cxx98-cu118-x86_64-linux/megablocks/grouped_gemm_util.py

@@ -4,7 +4,8 @@ import warnings
 
 _grouped_gemm_is_available: bool = False
 try:
-    import grouped_gemm
+    # import grouped_gemm
+    pass
     _grouped_gemm_is_available = True
 except ImportError as error:
     warnings.warn('Grouped GEMM not available.')
@@ -22,5 +23,9 @@ def assert_grouped_gemm_is_available():
     assert _grouped_gemm_is_available, msg
 
 
-backend = grouped_gemm.backend if grouped_gemm_is_available() else None
-ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+# backend = grouped_gemm.backend if grouped_gemm_is_available() else None
+# ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+
+
+from .grouped_gemm import backend as ops
+from .grouped_gemm import ops as backend

build/torch26-cxx98-cu118-x86_64-linux/megablocks/layers/__init__.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # from megablocks.layers.dmoe import dMoE
-from megablocks.layers.moe import MoE
+from .moe import MoE
 
 __all__ = [
     'MoE',

build/torch26-cxx98-cu124-x86_64-linux/megablocks/__init__.py

@@ -5,11 +5,15 @@ import torch
 
 from ._ops import ops
 
-from megablocks.layers.arguments import Arguments
-from megablocks.layers.dmoe import ParallelDroplessMLP, dMoE
-from megablocks.layers.glu import SparseGLU
-from megablocks.layers.mlp import MLP, SparseMLP
-from megablocks.layers.moe import MoE, ParallelMLP, get_load_balancing_loss
+from .grouped_gemm import backend as gg_backend
+from .grouped_gemm import ops as gg_ops
+
+
+from .layers.arguments import Arguments
+from .layers.dmoe import ParallelDroplessMLP, dMoE
+from .layers.glu import SparseGLU
+from .layers.mlp import MLP, SparseMLP
+from .layers.moe import MoE, ParallelMLP, get_load_balancing_loss
 
 # This section contains the direct kernel exports (not inlcuded in the original code)
 def exclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor:

build/torch26-cxx98-cu124-x86_64-linux/megablocks/_megablocks_359242d.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ec45bdb77d89916e5c58e17ddb76148799e3cac07fa6f4e93bf9140d9b2039bb
-size 11788400

build/torch26-cxx98-cu124-x86_64-linux/megablocks/_megablocks_a585153_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb0c2b91105c2f32f590aaa9d90ae2d6b36834bae9b35fb55c4b4fc90da56bc3
+size 11857952

build/torch26-cxx98-cu124-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_359242d
-ops = torch.ops._megablocks_359242d
+from . import _megablocks_a585153_dirty
+ops = torch.ops._megablocks_a585153_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_359242d::{op_name}"
+    return f"_megablocks_a585153_dirty::{op_name}"

build/torch26-cxx98-cu124-x86_64-linux/megablocks/grouped_gemm/__init__.py

@@ -0,0 +1,2 @@
+from . import ops
+from . import backend

build/torch26-cxx98-cu124-x86_64-linux/megablocks/grouped_gemm/backend.py

@@ -0,0 +1,32 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+
+# # TODO(tgale): Wrap this in a try-block with better
+# # error message and instructions for building the
+# # c++ operations.
+# import grouped_gemm_backend as backend
+
+# We import the backend operations from the megablocks package as
+# grouped_gemm is vendored in megablocks in this repository.
+# from ... import _ops as backend
+from megablocks._ops import ops as backend  # type: ignore
+
+def _allocate_output(a, b, batch_sizes, trans_a, trans_b):
+    assert not (trans_a and trans_b)
+    assert batch_sizes.ndim == 1, "Expected 1d tensor for batch_sizes"
+    assert a.ndim == 2, "Expected 2d tensor for 'a'"
+    assert b.ndim == (2 if trans_a else 3)
+
+    shape = (
+        (batch_sizes.shape[0], a.shape[1], b.shape[1])
+        if trans_a else
+        (a.shape[0], (b.shape[1] if trans_b else b.shape[2]))
+    )
+    return torch.empty(*shape, device=a.device, dtype=a.dtype)
+
+def gmm(a, b, batch_sizes, trans_a=False, trans_b=False, c=None):
+    if c is None:
+        c = _allocate_output(a, b, batch_sizes, trans_a, trans_b)
+    backend.gmm(a, b, c, batch_sizes, trans_a, trans_b)
+    return c

build/torch26-cxx98-cu124-x86_64-linux/megablocks/grouped_gemm/ops.py

@@ -0,0 +1,33 @@
+from . import backend
+import torch
+
+
+class GroupedGemm(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, a, b, batch_sizes, trans_b):
+        ctx.save_for_backward(a, b, batch_sizes)
+        ctx.trans_b = trans_b
+        return backend.gmm(a, b, batch_sizes, trans_a=False, trans_b=trans_b)
+
+    @staticmethod
+    def backward(ctx, grad):
+        grad = grad.contiguous()
+        a, b, batch_sizes = ctx.saved_tensors
+        trans_b = ctx.trans_b
+
+        agrad = None
+        if ctx.needs_input_grad[0]:
+            agrad = backend.gmm(
+                grad, b, batch_sizes, trans_a=False, trans_b=not trans_b)
+
+        bgrad = None
+        if ctx.needs_input_grad[1]:
+            lhs, rhs = (grad, a) if trans_b else (a, grad)
+            bgrad = backend.gmm(
+                lhs, rhs, batch_sizes, trans_a=True, trans_b=False)
+        return agrad, bgrad, None, None
+
+
+def gmm(a, b, batch_sizes, trans_b=False):
+    return GroupedGemm.apply(a, b, batch_sizes, trans_b)

build/torch26-cxx98-cu124-x86_64-linux/megablocks/grouped_gemm_util.py

@@ -4,7 +4,8 @@ import warnings
 
 _grouped_gemm_is_available: bool = False
 try:
-    import grouped_gemm
+    # import grouped_gemm
+    pass
     _grouped_gemm_is_available = True
 except ImportError as error:
     warnings.warn('Grouped GEMM not available.')
@@ -22,5 +23,9 @@ def assert_grouped_gemm_is_available():
     assert _grouped_gemm_is_available, msg
 
 
-backend = grouped_gemm.backend if grouped_gemm_is_available() else None
-ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+# backend = grouped_gemm.backend if grouped_gemm_is_available() else None
+# ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+
+
+from .grouped_gemm import backend as ops
+from .grouped_gemm import ops as backend

build/torch26-cxx98-cu124-x86_64-linux/megablocks/layers/__init__.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # from megablocks.layers.dmoe import dMoE
-from megablocks.layers.moe import MoE
+from .moe import MoE
 
 __all__ = [
     'MoE',

build/torch26-cxx98-cu126-x86_64-linux/megablocks/__init__.py

@@ -5,11 +5,15 @@ import torch
 
 from ._ops import ops
 
-from megablocks.layers.arguments import Arguments
-from megablocks.layers.dmoe import ParallelDroplessMLP, dMoE
-from megablocks.layers.glu import SparseGLU
-from megablocks.layers.mlp import MLP, SparseMLP
-from megablocks.layers.moe import MoE, ParallelMLP, get_load_balancing_loss
+from .grouped_gemm import backend as gg_backend
+from .grouped_gemm import ops as gg_ops
+
+
+from .layers.arguments import Arguments
+from .layers.dmoe import ParallelDroplessMLP, dMoE
+from .layers.glu import SparseGLU
+from .layers.mlp import MLP, SparseMLP
+from .layers.moe import MoE, ParallelMLP, get_load_balancing_loss
 
 # This section contains the direct kernel exports (not inlcuded in the original code)
 def exclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor:

build/torch26-cxx98-cu126-x86_64-linux/megablocks/_megablocks_359242d.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:99859c18a9a4e7ec11e7fa805e2225644f8e5f51e2546b4525ddf8e939f48874
-size 11832392

build/torch26-cxx98-cu126-x86_64-linux/megablocks/_megablocks_a585153_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd817eed5069e786933346cb2bb5ab6f586878ae80647191932336dec3295c96
+size 11923704

build/torch26-cxx98-cu126-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_359242d
-ops = torch.ops._megablocks_359242d
+from . import _megablocks_a585153_dirty
+ops = torch.ops._megablocks_a585153_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_359242d::{op_name}"
+    return f"_megablocks_a585153_dirty::{op_name}"

build/torch26-cxx98-cu126-x86_64-linux/megablocks/grouped_gemm/__init__.py

@@ -0,0 +1,2 @@
+from . import ops
+from . import backend

build/torch26-cxx98-cu126-x86_64-linux/megablocks/grouped_gemm/backend.py

@@ -0,0 +1,32 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+
+# # TODO(tgale): Wrap this in a try-block with better
+# # error message and instructions for building the
+# # c++ operations.
+# import grouped_gemm_backend as backend
+
+# We import the backend operations from the megablocks package as
+# grouped_gemm is vendored in megablocks in this repository.
+# from ... import _ops as backend
+from megablocks._ops import ops as backend  # type: ignore
+
+def _allocate_output(a, b, batch_sizes, trans_a, trans_b):
+    assert not (trans_a and trans_b)
+    assert batch_sizes.ndim == 1, "Expected 1d tensor for batch_sizes"
+    assert a.ndim == 2, "Expected 2d tensor for 'a'"
+    assert b.ndim == (2 if trans_a else 3)
+
+    shape = (
+        (batch_sizes.shape[0], a.shape[1], b.shape[1])
+        if trans_a else
+        (a.shape[0], (b.shape[1] if trans_b else b.shape[2]))
+    )
+    return torch.empty(*shape, device=a.device, dtype=a.dtype)
+
+def gmm(a, b, batch_sizes, trans_a=False, trans_b=False, c=None):
+    if c is None:
+        c = _allocate_output(a, b, batch_sizes, trans_a, trans_b)
+    backend.gmm(a, b, c, batch_sizes, trans_a, trans_b)
+    return c

build/torch26-cxx98-cu126-x86_64-linux/megablocks/grouped_gemm/ops.py

@@ -0,0 +1,33 @@
+from . import backend
+import torch
+
+
+class GroupedGemm(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, a, b, batch_sizes, trans_b):
+        ctx.save_for_backward(a, b, batch_sizes)
+        ctx.trans_b = trans_b
+        return backend.gmm(a, b, batch_sizes, trans_a=False, trans_b=trans_b)
+
+    @staticmethod
+    def backward(ctx, grad):
+        grad = grad.contiguous()
+        a, b, batch_sizes = ctx.saved_tensors
+        trans_b = ctx.trans_b
+
+        agrad = None
+        if ctx.needs_input_grad[0]:
+            agrad = backend.gmm(
+                grad, b, batch_sizes, trans_a=False, trans_b=not trans_b)
+
+        bgrad = None
+        if ctx.needs_input_grad[1]:
+            lhs, rhs = (grad, a) if trans_b else (a, grad)
+            bgrad = backend.gmm(
+                lhs, rhs, batch_sizes, trans_a=True, trans_b=False)
+        return agrad, bgrad, None, None
+
+
+def gmm(a, b, batch_sizes, trans_b=False):
+    return GroupedGemm.apply(a, b, batch_sizes, trans_b)

build/torch26-cxx98-cu126-x86_64-linux/megablocks/grouped_gemm_util.py

@@ -4,7 +4,8 @@ import warnings
 
 _grouped_gemm_is_available: bool = False
 try:
-    import grouped_gemm
+    # import grouped_gemm
+    pass
     _grouped_gemm_is_available = True
 except ImportError as error:
     warnings.warn('Grouped GEMM not available.')
@@ -22,5 +23,9 @@ def assert_grouped_gemm_is_available():
     assert _grouped_gemm_is_available, msg
 
 
-backend = grouped_gemm.backend if grouped_gemm_is_available() else None
-ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+# backend = grouped_gemm.backend if grouped_gemm_is_available() else None
+# ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+
+
+from .grouped_gemm import backend as ops
+from .grouped_gemm import ops as backend