[Kernels] Add swiglu and geglu MLP across models

src/transformers/integrations/hub_kernels.py

@@ -98,9 +98,9 @@
     )
     from kernels import use_kernelized_func as _kernels_use_kernelized_func
 
-    def use_kernel_forward_from_hub(layer_name: str):
+    def use_kernel_forward_from_hub(layer_name: str, condition: Callable[["nn.Module"], bool] | None = None):
         if _kernels_enabled:
-            return _kernels_use_kernel_forward_from_hub(layer_name)
+            return _kernels_use_kernel_forward_from_hub(layer_name, condition)
         else:
             logger.warning_once(
                 f"kernels hub usage is disabled through the environment USE_HUB_KERNELS={_TRANSFORMERS_USE_HUB_KERNELS}"

src/transformers/models/afmoe/modeling_afmoe.py

@@ -119,6 +119,7 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class AfmoeMLP(nn.Module):
     def __init__(self, config, intermediate_size=None):
         super().__init__()

src/transformers/models/aimv2/modeling_aimv2.py

@@ -97,6 +97,7 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class Aimv2MLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/aria/modeling_aria.py

@@ -215,6 +215,7 @@ def forward(self, key_value_states: torch.Tensor, attn_mask: torch.Tensor | None
         return out
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class AriaSharedExpertsMLP(nn.Module):
     """
     Shared Expert MLP for shared experts.

src/transformers/models/axk1/modeling_axk1.py

@@ -126,6 +126,7 @@ def forward(self, x, position_ids):
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class AXK1MLP(nn.Module):
     def __init__(self, config, intermediate_size=None):
         super().__init__()

src/transformers/models/axk2/modeling_axk2.py

@@ -381,6 +381,7 @@ def forward(
         return final_hidden_states
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class AXK2MLP(nn.Module):
     def __init__(self, config, intermediate_size=None):
         super().__init__()

src/transformers/models/bamba/modeling_bamba.py

@@ -795,6 +795,7 @@ def forward(
         return contextualized_states
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class BambaMLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/bitnet/modular_bitnet.py

@@ -22,6 +22,7 @@
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
 from ...processing_utils import Unpack
 from ...utils import logging
+from ...utils.generic import no_inherit_decorator
 from ..gemma.modeling_gemma import GemmaMLP
 from ..llama.modeling_llama import (
     LlamaAttention,
@@ -42,6 +43,7 @@ class BitNetRMSNorm(LlamaRMSNorm):
     pass
 
 
+@no_inherit_decorator
 class BitNetMLP(GemmaMLP):
     def __init__(self, config: BitNetConfig):
         super().__init__(config)

src/transformers/models/cohere/modeling_cohere.py

@@ -34,6 +34,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub
 from ...masking_utils import create_causal_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import GradientCheckpointingLayer
@@ -119,6 +120,7 @@ def forward(self, x, position_ids):
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class CohereMLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/cohere2/modeling_cohere2.py

@@ -26,6 +26,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub
 from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
@@ -253,6 +254,7 @@ def forward(
         return attn_output, attn_weights
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class Cohere2MLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/cohere2_moe/modeling_cohere2_moe.py

@@ -82,6 +82,7 @@ def forward(self, hidden_states):
         return hidden_states.to(input_dtype)
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class Cohere2MoeMLP(nn.Module):
     def __init__(self, config: Cohere2MoeConfig, intermediate_size=None):
         super().__init__()

src/transformers/models/cohere_compass/modeling_cohere_compass.py

@@ -160,6 +160,7 @@ def apply_interleaved_mrope(self, freqs, mrope_section):
         return freqs_t
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class CohereCompassMLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/csm/modeling_csm.py

@@ -174,6 +174,7 @@ def forward(self, x, position_ids):
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class CsmMLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/cwm/modeling_cwm.py

@@ -249,6 +249,7 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class CwmMLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py

@@ -1118,6 +1118,7 @@ def forward(
         return attn_output, attn_weights
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class DeepseekOcr2TextMLP(nn.Module):
     def __init__(self, config: DeepseekOcr2TextConfig, hidden_size=None, intermediate_size=None):
         super().__init__()

src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py

@@ -40,7 +40,7 @@
 from ...modeling_utils import PreTrainedModel
 from ...processing_utils import Unpack
 from ...utils import TensorType, TransformersKwargs, auto_docstring, can_return_tuple, logging
-from ...utils.generic import merge_with_config_defaults
+from ...utils.generic import merge_with_config_defaults, no_inherit_decorator
 from ...utils.import_utils import requires
 from ...utils.output_capturing import capture_outputs
 from ..deepseek_v2.configuration_deepseek_v2 import DeepseekV2Config
@@ -783,6 +783,7 @@ def forward(self, pixel_values: torch.FloatTensor, **kwargs) -> BaseModelOutput:
         return BaseModelOutput(last_hidden_state=hidden_states)
 
 
+@no_inherit_decorator
 class DeepseekOcr2VisionMLP(Qwen2MLP):
     pass
 

src/transformers/models/deepseek_v2/modeling_deepseek_v2.py

@@ -137,6 +137,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return hidden_states
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class DeepseekV2MLP(nn.Module):
     def __init__(self, config: DeepseekV2Config, hidden_size=None, intermediate_size=None):
         super().__init__()

src/transformers/models/deepseek_v3/modeling_deepseek_v3.py

@@ -112,6 +112,7 @@ def forward(self, x, position_ids):
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class DeepseekV3MLP(nn.Module):
     def __init__(self, config, intermediate_size=None):
         super().__init__()

src/transformers/models/deepseek_v32/modeling_deepseek_v32.py

@@ -503,6 +503,7 @@ def forward(
         return attn_output, attn_weights
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class DeepseekV32MLP(nn.Module):
     def __init__(self, config, intermediate_size=None):
         super().__init__()

src/transformers/models/deepseek_v4/modeling_deepseek_v4.py

@@ -962,6 +962,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return (pre.unsqueeze(-1) * x).sum(dim=2).to(x.dtype)
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class DeepseekV4MLP(nn.Module):
     def __init__(self, config: DeepseekV4Config):
         super().__init__()

src/transformers/models/deimv2/modeling_deimv2.py

@@ -177,6 +177,7 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class Deimv2SwiGLUFFN(nn.Module):
     def __init__(self, config: Deimv2Config):
         super().__init__()

src/transformers/models/diffllama/modeling_diffllama.py

@@ -49,6 +49,7 @@
 from .configuration_diffllama import DiffLlamaConfig
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class DiffLlamaMLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/diffusion_gemma/configuration_diffusion_gemma.py

@@ -75,6 +75,7 @@ class DiffusionGemmaTextConfig(PreTrainedConfig):
         "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
         "norm": (["hidden_states"], ["hidden_states"]),
     }
+    attribute_map = {"hidden_act": "hidden_activation"}
 
     vocab_size: int = 262_144
     hidden_size: int = 2304

src/transformers/models/diffusion_gemma/modeling_diffusion_gemma.py

@@ -28,7 +28,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...configuration_utils import PreTrainedConfig
-from ...integrations import use_experts_implementation
+from ...integrations import use_experts_implementation, use_kernel_forward_from_hub
 from ...masking_utils import (
     ALL_MASK_ATTENTION_FUNCTIONS,
     bidirectional_mask_function,
@@ -497,6 +497,7 @@ def append_to_cache(
         return keys, values
 
 
+@use_kernel_forward_from_hub("GeGLUMLP", condition=lambda module: module.config.hidden_act == "gelu_pytorch_tanh")
 class DiffusionGemmaText4MLP(nn.Module):
     def __init__(self, config: DiffusionGemmaTextConfig, layer_idx: int):
         super().__init__()
@@ -506,7 +507,7 @@ def __init__(self, config: DiffusionGemmaTextConfig, layer_idx: int):
         self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_activation]
+        self.act_fn = ACT2FN[config.hidden_act]
 
     def forward(self, x):
         down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py

@@ -464,7 +464,7 @@ def __init__(self, config: DiffusionGemmaTextConfig, layer_idx: int):
         self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_activation]
+        self.act_fn = ACT2FN[config.hidden_act]
 
 
 class DiffusionGemmaTextRouter(Gemma4TextRouter):

src/transformers/models/dinov3_vit/modular_dinov3_vit.py

@@ -38,7 +38,7 @@
 from ...processing_utils import Unpack
 from ...pytorch_utils import compile_compatible_method_lru_cache
 from ...utils import TransformersKwargs, auto_docstring, logging
-from ...utils.generic import can_return_tuple, maybe_autocast, merge_with_config_defaults
+from ...utils.generic import can_return_tuple, maybe_autocast, merge_with_config_defaults, no_inherit_decorator
 from ...utils.output_capturing import capture_outputs
 from ..swin.modeling_swin import SwinDropPath
 from .configuration_dinov3_vit import DINOv3ViTConfig
@@ -300,6 +300,7 @@ class DINOv3ViTMLP(ArceeMLP):
     pass
 
 
+@no_inherit_decorator
 class DINOv3ViTGatedMLP(LlamaMLP):
     pass
 

src/transformers/models/doge/modeling_doge.py

@@ -371,6 +371,7 @@ def prepare_dynamic_mask(
         return attn_mask
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class DogeMLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/dots1/modeling_dots1.py

@@ -260,6 +260,7 @@ def forward(
         return attn_output, attn_weights
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class Dots1MLP(nn.Module):
     def __init__(self, config, intermediate_size=None):
         super().__init__()

src/transformers/models/emu3/modeling_emu3.py

@@ -216,6 +216,7 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class Emu3MLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/ernie4_5/modeling_ernie4_5.py

@@ -95,6 +95,7 @@ def forward(self, x, position_ids):
         return cos, sin
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class Ernie4_5MLP(nn.Module):
     def __init__(self, config: Ernie4_5Config):
         super().__init__()

src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py

@@ -63,6 +63,7 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class Ernie4_5_MoeMLP(nn.Module):
     def __init__(self, config, intermediate_size=None):
         super().__init__()

src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py

@@ -290,6 +290,7 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class Ernie4_5_VLMoeMLP(nn.Module):
     def __init__(self, config, intermediate_size=None):
         super().__init__()

src/transformers/models/esmc/modeling_esmc.py

@@ -110,6 +110,7 @@ def forward(self, x, position_ids):
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class EsmcMLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/eurobert/modeling_eurobert.py

@@ -200,6 +200,7 @@ def forward(
         return attn_output, attn_weights
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class EuroBertMLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/evolla/modeling_evolla.py

@@ -1083,6 +1083,7 @@ def forward(self, x, position_ids):
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class EvollaMLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/exaone4/modeling_exaone4.py

@@ -265,6 +265,7 @@ def forward(
         return attn_output, attn_weights
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class Exaone4MLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/exaone_moe/modeling_exaone_moe.py

@@ -208,6 +208,7 @@ def forward(
         return attn_output, attn_weights
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class ExaoneMoeMLP(nn.Module):
     def __init__(self, config, intermediate_size=None):
         super().__init__()

src/transformers/models/falcon_h1/modeling_falcon_h1.py

@@ -822,6 +822,7 @@ def forward(
         return contextualized_states
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class FalconH1MLP(nn.Module):
     def __init__(self, config: FalconH1Config):
         super().__init__()

src/transformers/models/flex_olmo/modeling_flex_olmo.py

@@ -117,6 +117,7 @@ def forward(self, x, position_ids):
         return cos, sin
 
 
+@use_kernel_forward_from_hub("SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu")
 class FlexOlmoMLP(nn.Module):
     def __init__(self, config):
         super().__init__()

src/transformers/models/gemma/modeling_gemma.py

@@ -81,6 +81,7 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.eps}"
 
 
+@use_kernel_forward_from_hub("GeGLUMLP", condition=lambda module: module.config.hidden_act == "gelu_pytorch_tanh")
 class GemmaMLP(nn.Module):
     def __init__(self, config):
         super().__init__()