ScatterMoE

Author: shawntan

src/transformers/integrations/hub_kernels.py

@@ -115,6 +115,12 @@ def use_kernel_forward_from_hub(layer_name: str):
                 )
             },
         },
+        "ScatterMoEGatedMLP": {
+            "cuda": {
+                Mode.TRAINING: LayerRepository(repo_id="shawntan/scattermoe", layer_name="ScatterMoEGatedMLP"),
+                Mode.INFERENCE: LayerRepository(repo_id="shawntan/scattermoe", layer_name="ScatterMoEGatedMLP"),
+            },
+        },
         "FastGELU": {
             "cuda": {
                 Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository(

src/transformers/models/granitemoe/modeling_granitemoe.py

@@ -221,6 +221,7 @@ def forward(self, hidden_states):
         return index_sorted_experts, batch_index, batch_gates, expert_size, logits
 
 
+@use_kernel_forward_from_hub("ScatterMoEGatedMLP")
 class GraniteMoeMoE(nn.Module):
     """
     A Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.

src/transformers/models/granitemoe/modular_granitemoe.py

@@ -20,6 +20,7 @@
 
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
+from ...integrations import use_kernel_forward_from_hub
 from ...masking_utils import create_causal_mask
 from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
 from ...modeling_utils import PreTrainedModel
@@ -49,6 +50,7 @@ class GraniteMoeTopKGating(JetMoeTopKGating):
     pass
 
 
+@use_kernel_forward_from_hub("ScatterMoEGatedMLP")
 class GraniteMoeMoE(nn.Module):
     """
     A Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.

src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py

@@ -1066,6 +1066,7 @@ def forward(self, hidden_states):
         return index_sorted_experts, batch_index, batch_gates, expert_size, logits
 
 
+@use_kernel_forward_from_hub("ScatterMoEGatedMLP")
 class GraniteMoeHybridMoE(nn.Module):
     """
     A Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.

src/transformers/models/granitemoeshared/modeling_granitemoeshared.py

@@ -207,6 +207,7 @@ def forward(self, hidden_states):
         return index_sorted_experts, batch_index, batch_gates, expert_size, logits
 
 
+@use_kernel_forward_from_hub("ScatterMoEGatedMLP")
 class GraniteMoeSharedMoE(nn.Module):
     """
     A Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.