Merge pull request #7 from huggingface/megablocks_moe

feat: add megablocks moe mlp kernel

Author: drbh

src/transformers/integrations/hub_kernels.py

@@ -23,9 +23,7 @@
         register_kernel_mapping,
         replace_kernel_forward_from_hub,
     )
-    from kernels import (
-        use_kernel_forward_from_hub as original_use_kernel_forward_from_hub,
-    )
+    from kernels import use_kernel_forward_from_hub
 
     _hub_kernels_available = True
 
@@ -56,44 +54,16 @@
                 layer_name="TritonLlamaMLP",
             )
         },
+        "MegaBlocksMoeMLP": {
+            "cuda": LayerRepository(
+                repo_id="kernels-community/megablocks",
+                layer_name="MegaBlocksMoeMLP",
+            )
+        },
     }
 
     register_kernel_mapping(_KERNEL_MAPPING)
 
-    def use_kernel_forward_from_hub(*args, **kwargs):
-        """
-        Expands `kernels`' `use_kernel_forward_from_hub` to NOT use a kernel at compile time. This should be removed
-        when `kernels` supports `torch.compile`.
-
-        If the layer has a `config` attribute, we can also set `config.disable_custom_kernels = True` to disable the
-        kernel.
-        """
-
-        def decorator_with_compile_path(cls):
-            # Keeps a reference to the original forward method
-            original_forward = cls.forward
-
-            # Applies the original decorator
-            decorator = original_use_kernel_forward_from_hub(*args, **kwargs)
-            cls = decorator(cls)
-
-            # Replaces the kernel forward with a compile-friendly version
-            kernel_forward = cls.forward
-
-            def forward_with_compile_path(*forward_args, **forward_kwargs):
-                disable_custom_kernels = hasattr(cls, "config") and getattr(cls.config, "disable_custom_kernels", None)
-                if is_torchdynamo_compiling() or disable_custom_kernels:
-                    return original_forward(*forward_args, **forward_kwargs)
-                else:
-                    return kernel_forward(*forward_args, **forward_kwargs)
-
-            cls.forward = forward_with_compile_path
-
-            return cls
-
-        return decorator_with_compile_path
-
-
 except ImportError:
     # Stub to make decorators int transformers work when `kernels`
     # is not installed.

src/transformers/models/openai_moe/modeling_openai_moe.py

@@ -123,6 +123,7 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
         return next_states
 
 
+@use_kernel_forward_from_hub("MegaBlocksMoeMLP")
 class OpenAIMoeMLP(nn.Module):
     def __init__(self, config):
         super().__init__()