Remove MLP class.

Author: shawntan

scattermoe/README.md

@@ -1,3 +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/torch-ext/scattermoe/__init__.py

@@ -1,7 +1,6 @@
 from .parallel_experts import flatten_sort_count, parallel_linear, ParallelExperts
 from . import parallel_experts
 from . import kernels
-from . import mlp
 from . import layers
 
 __all__ = [
@@ -10,6 +9,5 @@
     "ParallelExperts",
     "parallel_experts",
     "kernels",
-    "mlp",
     "layers"
 ]

scattermoe/torch-ext/scattermoe/layers.py

@@ -48,5 +48,5 @@ def forward(self, layer_input):
             gates=routing_weights
         )
         layer_output = layer_output.view(bsz, length, emb_size)
-        return layer_output
+        return layer_output, router_logits