scattermoe: add

scattermoe/README.md

@@ -0,0 +1,5 @@
+# ScatterMoE
+
+This is the [ScatterMoE](https://arxiv.org/abs/2403.08245) written in Triton
+
+The main repository for ScatterMoE is [here](https://github.com/shawntan/scattermoe)

scattermoe/build.toml

@@ -0,0 +1,3 @@
+[general]
+name = "scattermoe"
+universal = true

scattermoe/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for scattermoe kernels";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernel-builder";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

scattermoe/torch-ext/scattermoe/__init__.py

@@ -0,0 +1,13 @@
+from .parallel_experts import flatten_sort_count, parallel_linear, ParallelExperts
+from . import parallel_experts
+from . import kernels
+from . import layers
+
+__all__ = [
+    "flatten_sort_count",
+    "parallel_linear",
+    "ParallelExperts",
+    "parallel_experts",
+    "kernels",
+    "layers"
+]

scattermoe/torch-ext/scattermoe/kernels/__init__.py

@@ -0,0 +1,3 @@
+from . import ops
+
+__all__ = ["ops"]