From c82c2efeb928be579bd4140a4e150806e0c6394e Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <ilyas.moutawwakil@gmail.com>
Date: Wed, 8 Apr 2026 12:55:20 +0000
Subject: [PATCH 01/13] added sonic moe

---
 src/transformers/integrations/moe.py      |   3 +
 src/transformers/integrations/sonicmoe.py | 126 ++++++++++++++++++++++
 2 files changed, 129 insertions(+)
 create mode 100644 src/transformers/integrations/sonicmoe.py

diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py
index cd584278a869..6c6893e49479 100644
--- a/src/transformers/integrations/moe.py
+++ b/src/transformers/integrations/moe.py
@@ -18,6 +18,7 @@
 from ..utils import logging
 from ..utils.generic import GeneralInterface
 from ..utils.import_utils import is_torch_available, is_torch_less_or_equal, is_torchdynamo_compiling
+from .sonicmoe import sonicmoe_experts_forward
 
 
 if is_torch_available():
@@ -26,6 +27,7 @@
 
 logger = logging.get_logger(__name__)
 
+
 # Examples of experts class with its eager mm implementation
 # class Experts(torch.nn.Module):
 #     """Collection of expert weights stored as 3D tensors."""
@@ -433,6 +435,7 @@ class ExpertsInterface(GeneralInterface):
     """Interface for registering custom experts forward functions."""
 
     _global_mapping = {
+        "sonicmoe": sonicmoe_experts_forward,
         "batched_mm": batched_mm_experts_forward,
         "grouped_mm": grouped_mm_experts_forward,
     }
diff --git a/src/transformers/integrations/sonicmoe.py b/src/transformers/integrations/sonicmoe.py
new file mode 100644
index 000000000000..0fcd2ee088a1
--- /dev/null
+++ b/src/transformers/integrations/sonicmoe.py
@@ -0,0 +1,126 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SonicMoE integration: fused MoE using CuteDSL kernels from `kernels-community/sonic-moe`.
+
+Provides `sonicmoe_experts_forward` registered as "sonicmoe" in the ExpertsInterface.
+Requirements: CUDA, `kernels`, `nvidia-cutlass-dsl`, has_gate=True, is_transposed=False.
+"""
+
+import os
+
+import torch
+
+from ..utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+_sonicmoe_module = None
+
+# Map activation function names from HF config to SonicMoE epilogue names
+ACT_MAP = {"silu": "swiglu", "gelu": "geglu", "relu": "reglu"}
+
+
+def _load_sonicmoe():
+    """Lazy-load the SonicMoE kernel from hub."""
+    global _sonicmoe_module
+    if _sonicmoe_module is None:
+        if "CUTE_DSL_ARCH" not in os.environ:
+            major, minor = torch.cuda.get_device_capability()
+            os.environ["CUTE_DSL_ARCH"] = f"sm_{major}{minor}a" if major >= 9 else f"sm_{major}{minor}"
+
+        from kernels import get_kernel
+
+        _sonicmoe_module = get_kernel("adarshxs/sonic-moe", revision="v1")
+
+    return _sonicmoe_module
+
+
+def sonicmoe_experts_forward(
+    self: torch.nn.Module,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    if self.is_transposed:
+        raise ValueError("sonicmoe requires non-transposed weights (is_transposed=False)")
+    if not self.has_gate:
+        raise ValueError("sonicmoe requires gated experts (has_gate=True)")
+    if hidden_states.device.type != "cuda":
+        raise ValueError("sonicmoe requires CUDA device")
+
+    kernel = _load_sonicmoe()
+    moe_general_routing = getattr(kernel, "moe_general_routing_inputs", None)
+    if moe_general_routing is None:
+        raise ImportError(
+            "moe_general_routing_inputs not found in adarshxs/sonic-moe. "
+            "Make sure you have the `kernels` package and `nvidia-cutlass-dsl` installed."
+        )
+
+    ActivationType = getattr(getattr(kernel, "enums", None), "ActivationType", None)
+    if ActivationType is None:
+        raise ImportError(
+            "ActivationType enum not found in adarshxs/sonic-moe. "
+            "Make sure you have the `kernels` package and `nvidia-cutlass-dsl` installed."
+        )
+
+    device = hidden_states.device
+    num_experts = self.num_experts
+    num_top_k = top_k_index.size(-1)
+    num_tokens = hidden_states.size(0)
+    grad_enabled = torch.is_grad_enabled()
+    cuda_device_stream = torch.cuda.current_stream(device).cuda_stream
+
+    # Flatten — token_indices must be int32, sorted ascending (required by sonic-moe)
+    token_idx = (
+        torch.arange(num_tokens, device=device, dtype=torch.int32).unsqueeze(1).expand(-1, num_top_k).reshape(-1)
+    )
+    router_scores = top_k_weights.reshape(-1).to(hidden_states.dtype)
+    expert_ids = top_k_index.reshape(-1).to(torch.int32)
+
+    # Map activation function
+    act_name = getattr(self, "act_fn", "silu").lower()
+    activation_type = getattr(ActivationType, ACT_MAP.get(act_name, "swiglu").upper(), ActivationType.SWIGLU)
+
+    # SonicMoE's CUTLASS epilogue reads accumulators as interleaved pairs: act(D[2i]) * D[2i+1].
+    # HuggingFace stores concatenated [gate; up]. Interleave and cache on first call.
+    if not hasattr(self, "is_interleaved"):
+        I = self.gate_up_proj.size(1) // 2
+        self.gate_up_proj = torch.stack([self.gate_up_proj[:, :I, :], self.gate_up_proj[:, I:, :]], dim=2).reshape(
+            self.gate_up_proj.shape
+        )
+        self.is_interleaved = True
+
+    w1 = self.gate_up_proj.permute(1, 2, 0)  # (2*I, H, E)
+    w2 = self.down_proj.permute(1, 2, 0)  # (I, H, E)
+    b1 = self.gate_up_proj_bias if self.has_bias else None
+    b2 = self.down_proj_bias if self.has_bias else None
+
+    output, _ = moe_general_routing(
+        hidden_states,
+        router_scores,
+        token_idx,
+        expert_ids,
+        w1,
+        b1,
+        w2,
+        b2,
+        num_experts,
+        cuda_device_stream,
+        activation_type,
+        grad_enabled,
+    )
+
+    return output

From ebf3bdc6633d4e0eb590e2437d4aeca4c92abde5 Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <ilyas.moutawwakil@gmail.com>
Date: Tue, 14 Apr 2026 14:24:52 +0000
Subject: [PATCH 02/13] use lazy_load_kernel

---
 src/transformers/integrations/hub_kernels.py |  1 +
 src/transformers/integrations/sonicmoe.py    | 39 ++++++--------------
 2 files changed, 13 insertions(+), 27 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index 39c1448af02b..4fbdebe04607 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -286,6 +286,7 @@ def register_kernel_mapping_transformers(*args, **kwargs):
     "falcon_mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1},
     "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 1},
     "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 1},
+    "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "version": 1},
 }
 
 _KERNEL_MODULE_MAPPING: dict[str, ModuleType | None] = {}
diff --git a/src/transformers/integrations/sonicmoe.py b/src/transformers/integrations/sonicmoe.py
index 0fcd2ee088a1..7159973dbf70 100644
--- a/src/transformers/integrations/sonicmoe.py
+++ b/src/transformers/integrations/sonicmoe.py
@@ -23,29 +23,22 @@
 import torch
 
 from ..utils import logging
+from .hub_kernels import lazy_load_kernel
 
 
 logger = logging.get_logger(__name__)
 
-_sonicmoe_module = None
-
 # Map activation function names from HF config to SonicMoE epilogue names
 ACT_MAP = {"silu": "swiglu", "gelu": "geglu", "relu": "reglu"}
 
 
 def _load_sonicmoe():
-    """Lazy-load the SonicMoE kernel from hub."""
-    global _sonicmoe_module
-    if _sonicmoe_module is None:
-        if "CUTE_DSL_ARCH" not in os.environ:
-            major, minor = torch.cuda.get_device_capability()
-            os.environ["CUTE_DSL_ARCH"] = f"sm_{major}{minor}a" if major >= 9 else f"sm_{major}{minor}"
-
-        from kernels import get_kernel
-
-        _sonicmoe_module = get_kernel("adarshxs/sonic-moe", revision="v1")
+    """Lazy-load the SonicMoE kernel (hub or local install)."""
+    if "CUTE_DSL_ARCH" not in os.environ:
+        major, minor = torch.cuda.get_device_capability()
+        os.environ["CUTE_DSL_ARCH"] = f"sm_{major}{minor}a" if major >= 9 else f"sm_{major}{minor}"
 
-    return _sonicmoe_module
+    return lazy_load_kernel("sonic-moe")
 
 
 def sonicmoe_experts_forward(
@@ -65,14 +58,14 @@ def sonicmoe_experts_forward(
     moe_general_routing = getattr(kernel, "moe_general_routing_inputs", None)
     if moe_general_routing is None:
         raise ImportError(
-            "moe_general_routing_inputs not found in adarshxs/sonic-moe. "
+            "moe_general_routing_inputs not found in kernels-community/sonic-moe. "
             "Make sure you have the `kernels` package and `nvidia-cutlass-dsl` installed."
         )
 
     ActivationType = getattr(getattr(kernel, "enums", None), "ActivationType", None)
     if ActivationType is None:
         raise ImportError(
-            "ActivationType enum not found in adarshxs/sonic-moe. "
+            "ActivationType enum not found in kernels-community/sonic-moe. "
             "Make sure you have the `kernels` package and `nvidia-cutlass-dsl` installed."
         )
 
@@ -91,18 +84,9 @@ def sonicmoe_experts_forward(
     expert_ids = top_k_index.reshape(-1).to(torch.int32)
 
     # Map activation function
-    act_name = getattr(self, "act_fn", "silu").lower()
+    act_name = getattr(self.config, "hidden_act", "silu").lower()
     activation_type = getattr(ActivationType, ACT_MAP.get(act_name, "swiglu").upper(), ActivationType.SWIGLU)
 
-    # SonicMoE's CUTLASS epilogue reads accumulators as interleaved pairs: act(D[2i]) * D[2i+1].
-    # HuggingFace stores concatenated [gate; up]. Interleave and cache on first call.
-    if not hasattr(self, "is_interleaved"):
-        I = self.gate_up_proj.size(1) // 2
-        self.gate_up_proj = torch.stack([self.gate_up_proj[:, :I, :], self.gate_up_proj[:, I:, :]], dim=2).reshape(
-            self.gate_up_proj.shape
-        )
-        self.is_interleaved = True
-
     w1 = self.gate_up_proj.permute(1, 2, 0)  # (2*I, H, E)
     w2 = self.down_proj.permute(1, 2, 0)  # (I, H, E)
     b1 = self.gate_up_proj_bias if self.has_bias else None
@@ -119,8 +103,9 @@ def sonicmoe_experts_forward(
         b2,
         num_experts,
         cuda_device_stream,
-        activation_type,
-        grad_enabled,
+        activation_type=activation_type,
+        is_inference_mode_enabled=not grad_enabled,
+        is_concatenated_gate_up=True,
     )
 
     return output

From 7dd8cd07cf851811a85ead126f83a250e217a1cd Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <ilyas.moutawwakil@gmail.com>
Date: Tue, 14 Apr 2026 14:35:00 +0000
Subject: [PATCH 03/13] style

---
 src/transformers/integrations/moe.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py
index 9a76a61c8b49..da32ae23d4f0 100644
--- a/src/transformers/integrations/moe.py
+++ b/src/transformers/integrations/moe.py
@@ -26,7 +26,6 @@
 from .sonicmoe import sonicmoe_experts_forward
 
 
-
 if is_torch_available():
     import torch
 

From 0a7b8c01b06f5299234a78e3abc556664d9316b6 Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <ilyas.moutawwakil@gmail.com>
Date: Fri, 17 Apr 2026 10:15:13 +0000
Subject: [PATCH 04/13] use concatenated revision

---
 src/transformers/integrations/hub_kernels.py |  2 +-
 src/transformers/integrations/sonicmoe.py    | 40 ++++++--------------
 2 files changed, 12 insertions(+), 30 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index f3471dfed69a..3259bf6e0587 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -286,7 +286,7 @@ def register_kernel_mapping_transformers(*args, **kwargs):
     "falcon_mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1},
     "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 1},
     "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 1},
-    "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "version": 1},
+    "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "concatenated-gate-up"},
 }
 
 _KERNEL_MODULE_MAPPING: dict[str, ModuleType | None] = {}
diff --git a/src/transformers/integrations/sonicmoe.py b/src/transformers/integrations/sonicmoe.py
index 7159973dbf70..159e1518b570 100644
--- a/src/transformers/integrations/sonicmoe.py
+++ b/src/transformers/integrations/sonicmoe.py
@@ -18,8 +18,6 @@
 Requirements: CUDA, `kernels`, `nvidia-cutlass-dsl`, has_gate=True, is_transposed=False.
 """
 
-import os
-
 import torch
 
 from ..utils import logging
@@ -32,15 +30,6 @@
 ACT_MAP = {"silu": "swiglu", "gelu": "geglu", "relu": "reglu"}
 
 
-def _load_sonicmoe():
-    """Lazy-load the SonicMoE kernel (hub or local install)."""
-    if "CUTE_DSL_ARCH" not in os.environ:
-        major, minor = torch.cuda.get_device_capability()
-        os.environ["CUTE_DSL_ARCH"] = f"sm_{major}{minor}a" if major >= 9 else f"sm_{major}{minor}"
-
-    return lazy_load_kernel("sonic-moe")
-
-
 def sonicmoe_experts_forward(
     self: torch.nn.Module,
     hidden_states: torch.Tensor,
@@ -54,18 +43,12 @@ def sonicmoe_experts_forward(
     if hidden_states.device.type != "cuda":
         raise ValueError("sonicmoe requires CUDA device")
 
-    kernel = _load_sonicmoe()
-    moe_general_routing = getattr(kernel, "moe_general_routing_inputs", None)
-    if moe_general_routing is None:
+    sonic_moe = lazy_load_kernel("sonic-moe")
+    moe_general_routing = getattr(sonic_moe, "moe_general_routing_inputs", None)
+    ActivationType = getattr(getattr(sonic_moe, "enums", None), "ActivationType", None)
+    if moe_general_routing is None or ActivationType is None:
         raise ImportError(
-            "moe_general_routing_inputs not found in kernels-community/sonic-moe. "
-            "Make sure you have the `kernels` package and `nvidia-cutlass-dsl` installed."
-        )
-
-    ActivationType = getattr(getattr(kernel, "enums", None), "ActivationType", None)
-    if ActivationType is None:
-        raise ImportError(
-            "ActivationType enum not found in kernels-community/sonic-moe. "
+            "moe_general_routing_inputs function or ActivationType enum not found in kernels-community/sonic-moe. "
             "Make sure you have the `kernels` package and `nvidia-cutlass-dsl` installed."
         )
 
@@ -74,19 +57,18 @@ def sonicmoe_experts_forward(
     num_top_k = top_k_index.size(-1)
     num_tokens = hidden_states.size(0)
     grad_enabled = torch.is_grad_enabled()
-    cuda_device_stream = torch.cuda.current_stream(device).cuda_stream
+    stream_id = torch.cuda.current_stream(device).cuda_stream
 
     # Flatten — token_indices must be int32, sorted ascending (required by sonic-moe)
-    token_idx = (
-        torch.arange(num_tokens, device=device, dtype=torch.int32).unsqueeze(1).expand(-1, num_top_k).reshape(-1)
-    )
+    token_idx = torch.arange(num_tokens, device=device).unsqueeze(1).expand(-1, num_top_k).reshape(-1).int()
     router_scores = top_k_weights.reshape(-1).to(hidden_states.dtype)
-    expert_ids = top_k_index.reshape(-1).to(torch.int32)
+    expert_ids = top_k_index.reshape(-1).int()
 
     # Map activation function
     act_name = getattr(self.config, "hidden_act", "silu").lower()
     activation_type = getattr(ActivationType, ACT_MAP.get(act_name, "swiglu").upper(), ActivationType.SWIGLU)
 
+    # Permute weights to (E, H, I) as expected by sonic-moe (E=num_experts, H=hidden_size, I=intermediate_size)
     w1 = self.gate_up_proj.permute(1, 2, 0)  # (2*I, H, E)
     w2 = self.down_proj.permute(1, 2, 0)  # (I, H, E)
     b1 = self.gate_up_proj_bias if self.has_bias else None
@@ -101,8 +83,8 @@ def sonicmoe_experts_forward(
         b1,
         w2,
         b2,
-        num_experts,
-        cuda_device_stream,
+        E=num_experts,
+        stream_id=stream_id,
         activation_type=activation_type,
         is_inference_mode_enabled=not grad_enabled,
         is_concatenated_gate_up=True,

From 71d7e90084f93a39a742e6c401217e01d6f07135 Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <ilyas.moutawwakil@gmail.com>
Date: Thu, 23 Apr 2026 07:34:41 +0000
Subject: [PATCH 05/13] final touches

---
 src/transformers/integrations/hub_kernels.py  |  2 +-
 src/transformers/integrations/moe.py          | 10 +++-
 src/transformers/integrations/sonicmoe.py     | 60 ++++++++++++++-----
 src/transformers/modeling_utils.py            |  6 +-
 .../models/gpt_oss/modeling_gpt_oss.py        |  2 +-
 .../models/gpt_oss/modular_gpt_oss.py         |  2 +-
 6 files changed, 61 insertions(+), 21 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index 3259bf6e0587..f3471dfed69a 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -286,7 +286,7 @@ def register_kernel_mapping_transformers(*args, **kwargs):
     "falcon_mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1},
     "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 1},
     "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 1},
-    "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "concatenated-gate-up"},
+    "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "version": 1},
 }
 
 _KERNEL_MODULE_MAPPING: dict[str, ModuleType | None] = {}
diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py
index da32ae23d4f0..23d0cbdeac4f 100644
--- a/src/transformers/integrations/moe.py
+++ b/src/transformers/integrations/moe.py
@@ -495,6 +495,7 @@ def use_experts_implementation(
     experts_class: type[torch.nn.Module] | None = None,
     *,
     experts_interface: ExpertsInterface = ALL_EXPERTS_FUNCTIONS,
+    is_concatenated_layout: bool = True,
     is_transposed: bool = False,
     has_bias: bool = False,
     has_gate: bool = True,
@@ -506,10 +507,16 @@ def use_experts_implementation(
             The experts class to modify. If not provided, returns a decorator that can be applied to the class.
         experts_interface (`ExpertsInterface`, *optional*, defaults to `ALL_EXPERTS_FUNCTIONS`):
             The experts interface to use for dispatching the forward method.
+        is_concatenated_layout (`bool`, *optional*, defaults to `True`):
+            Whether the expert weights are stored in concatenated layout [gate;up]
+            or interleaved layout [gate0, up0, gate1, up1, ...].
         is_transposed (`bool`, *optional*, defaults to `False`):
             Whether the expert weights are stored in transposed format.
         has_bias (`bool`, *optional*, defaults to `False`):
-            Whether the expert layers include bias terms.
+            Whether the expert layers include bias terms or not.
+        has_gate (`bool`, *optional*, defaults to `True`):
+            Whether the experts use a gating mechanism or not.
+            Whether it has gate_up_proj weights or just up_proj weights.
 
     Returns:
         `type[torch.nn.Module]`: The modified experts class.
@@ -526,6 +533,7 @@ def __init__(self, config, *args, **kwargs):
             self.has_gate = has_gate
             self.has_bias = has_bias
             self.is_transposed = is_transposed
+            self.is_concatenated_layout = is_concatenated_layout
 
         @wraps(original_forward)
         def forward(self, *args, **kwargs):
diff --git a/src/transformers/integrations/sonicmoe.py b/src/transformers/integrations/sonicmoe.py
index 159e1518b570..7250ccd31640 100644
--- a/src/transformers/integrations/sonicmoe.py
+++ b/src/transformers/integrations/sonicmoe.py
@@ -18,6 +18,8 @@
 Requirements: CUDA, `kernels`, `nvidia-cutlass-dsl`, has_gate=True, is_transposed=False.
 """
 
+import functools
+
 import torch
 
 from ..utils import logging
@@ -30,6 +32,44 @@
 ACT_MAP = {"silu": "swiglu", "gelu": "geglu", "relu": "reglu"}
 
 
+@functools.cache
+def _load_sonic_kernel():
+    """
+    Load sonic-moe once and return its required symbols.
+
+    Raises:
+        ImportError if the kernel or required symbols are not found.
+
+    Returns:
+        Tuple of (ActivationType, moe_general_routing_inputs function) from the sonic-moe kernel.
+    """
+
+    kernel = lazy_load_kernel("sonic-moe")
+    if kernel is None:
+        raise ImportError(
+            "sonic-moe kernel not found. Make sure you have the `kernels` and `nvidia-cutlass-dsl` packages installed."
+        )
+
+    ActivationType = getattr(getattr(kernel, "enums", None), "ActivationType", None)
+    moe_general_routing_inputs = getattr(kernel, "moe_general_routing_inputs", None)
+
+    missing = [
+        name
+        for name, attr in [
+            ("enums.ActivationType", ActivationType),
+            ("moe_general_routing_inputs", moe_general_routing_inputs),
+        ]
+        if attr is None
+    ]
+    if missing:
+        raise ImportError(
+            f"sonic-moe kernel is missing required symbols: {', '.join(missing)}. "
+            "Make sure you have the `kernels` package and `nvidia-cutlass-dsl` installed."
+        )
+
+    return ActivationType, moe_general_routing_inputs
+
+
 def sonicmoe_experts_forward(
     self: torch.nn.Module,
     hidden_states: torch.Tensor,
@@ -43,21 +83,13 @@ def sonicmoe_experts_forward(
     if hidden_states.device.type != "cuda":
         raise ValueError("sonicmoe requires CUDA device")
 
-    sonic_moe = lazy_load_kernel("sonic-moe")
-    moe_general_routing = getattr(sonic_moe, "moe_general_routing_inputs", None)
-    ActivationType = getattr(getattr(sonic_moe, "enums", None), "ActivationType", None)
-    if moe_general_routing is None or ActivationType is None:
-        raise ImportError(
-            "moe_general_routing_inputs function or ActivationType enum not found in kernels-community/sonic-moe. "
-            "Make sure you have the `kernels` package and `nvidia-cutlass-dsl` installed."
-        )
+    ActivationType, moe_general_routing_inputs = _load_sonic_kernel()
 
     device = hidden_states.device
-    num_experts = self.num_experts
     num_top_k = top_k_index.size(-1)
     num_tokens = hidden_states.size(0)
     grad_enabled = torch.is_grad_enabled()
-    stream_id = torch.cuda.current_stream(device).cuda_stream
+    stream_id = torch.cuda.current_stream(device).cuda_stream  # this breaks torch.compile on some versions
 
     # Flatten — token_indices must be int32, sorted ascending (required by sonic-moe)
     token_idx = torch.arange(num_tokens, device=device).unsqueeze(1).expand(-1, num_top_k).reshape(-1).int()
@@ -68,13 +100,13 @@ def sonicmoe_experts_forward(
     act_name = getattr(self.config, "hidden_act", "silu").lower()
     activation_type = getattr(ActivationType, ACT_MAP.get(act_name, "swiglu").upper(), ActivationType.SWIGLU)
 
-    # Permute weights to (E, H, I) as expected by sonic-moe (E=num_experts, H=hidden_size, I=intermediate_size)
+    # Permute weights as expected by sonic-moe (E=num_experts, H=hidden_size, I=intermediate_size)
     w1 = self.gate_up_proj.permute(1, 2, 0)  # (2*I, H, E)
     w2 = self.down_proj.permute(1, 2, 0)  # (I, H, E)
     b1 = self.gate_up_proj_bias if self.has_bias else None
     b2 = self.down_proj_bias if self.has_bias else None
 
-    output, _ = moe_general_routing(
+    output, _ = moe_general_routing_inputs(
         hidden_states,
         router_scores,
         token_idx,
@@ -83,11 +115,11 @@ def sonicmoe_experts_forward(
         b1,
         w2,
         b2,
-        E=num_experts,
+        E=self.num_experts,
         stream_id=stream_id,
         activation_type=activation_type,
         is_inference_mode_enabled=not grad_enabled,
-        is_concatenated_gate_up=True,
+        concat_layout=self.is_concatenated_layout,
     )
 
     return output
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 2e98863a762d..1aebda423f06 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -1925,11 +1925,11 @@ def get_correct_attn_implementation(self, requested_attention: str | None, is_in
 
     def get_correct_experts_implementation(self, requested_experts: str | None) -> str:
         applicable_experts = "grouped_mm" if requested_experts is None else requested_experts
-        if applicable_experts not in ["eager", "grouped_mm", "batched_mm", "deepgemm"]:
+        if applicable_experts not in ["eager", "grouped_mm", "batched_mm", "deepgemm", "sonicmoe"]:
             message = (
                 f'Specified `experts_implementation="{applicable_experts}"` is not supported. The only possible arguments are '
-                '`experts_implementation="eager"`, `"experts_implementation=grouped_mm"`, `"experts_implementation=batched_mm"` '
-                'and `"experts_implementation=deepgemm"`.'
+                '`experts_implementation="eager"`, `"experts_implementation=grouped_mm"`, `"experts_implementation=batched_mm"`, '
+                '`experts_implementation=deepgemm`, `"experts_implementation=sonicmoe"`'
             )
             raise ValueError(message)
 
diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py
index 18f31ea90379..5f22a9654e94 100644
--- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py
+++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py
@@ -65,7 +65,7 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
-@use_experts_implementation(is_transposed=True, has_bias=True)
+@use_experts_implementation(is_concatenated_layout=False, is_transposed=True, has_bias=True)
 class GptOssExperts(nn.Module):
     def __init__(self, config):
         super().__init__()
diff --git a/src/transformers/models/gpt_oss/modular_gpt_oss.py b/src/transformers/models/gpt_oss/modular_gpt_oss.py
index 687e8864efeb..9c357e6e69ca 100644
--- a/src/transformers/models/gpt_oss/modular_gpt_oss.py
+++ b/src/transformers/models/gpt_oss/modular_gpt_oss.py
@@ -62,7 +62,7 @@ def forward(self, hidden_states):
         return (self.weight * hidden_states).to(input_dtype)  # main diff with Llama
 
 
-@use_experts_implementation(is_transposed=True, has_bias=True)
+@use_experts_implementation(is_concatenated_layout=False, is_transposed=True, has_bias=True)
 class GptOssExperts(nn.Module):
     def __init__(self, config):
         super().__init__()

From b078cb13495b0e82cb0e731692cc248713a4264a Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <ilyas.moutawwakil@gmail.com>
Date: Thu, 23 Apr 2026 11:45:31 +0000
Subject: [PATCH 06/13] fix

---
 src/transformers/modeling_utils.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 6699f83a64ed..8afe05b0ad60 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -1971,13 +1971,12 @@ def get_correct_attn_implementation(self, requested_attention: str | None, is_in
 
     def get_correct_experts_implementation(self, requested_experts: str | None) -> str:
         applicable_experts = "grouped_mm" if requested_experts is None else requested_experts
-<<<<<<< sonic-moe
         if applicable_experts not in ["eager", "grouped_mm", "batched_mm", "deepgemm", "sonicmoe"]:
             message = (
                 f'Specified `experts_implementation="{applicable_experts}"` is not supported. The only possible arguments are '
                 '`experts_implementation="eager"`, `"experts_implementation=grouped_mm"`, `"experts_implementation=batched_mm"`, '
                 '`experts_implementation=deepgemm`, `"experts_implementation=sonicmoe"`'
-=======
+            )
         base_experts_fns = ["eager"] + list(set(ALL_EXPERTS_FUNCTIONS.keys()) | set(ALL_FP8_EXPERTS_FUNCTIONS.keys()))
         valid_experts_str_list = [f'`experts_implementation="{fn}"`' for fn in base_experts_fns]
         valid_experts_str_list[-1] = "and " + valid_experts_str_list[-1]
@@ -1986,7 +1985,6 @@ def get_correct_experts_implementation(self, requested_experts: str | None) -> s
             message = (
                 f'Specified `experts_implementation="{applicable_experts}"` is not supported. The only possible arguments are '
                 f"{valid_experts_str}."
->>>>>>> main
             )
             raise ValueError(message)
 

From 134cf630512fb15e797b1fa8d724aa0d4b92c82a Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <ilyas.moutawwakil@gmail.com>
Date: Thu, 23 Apr 2026 11:50:34 +0000
Subject: [PATCH 07/13] merge conflict

---
 src/transformers/modeling_utils.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 8afe05b0ad60..d58c9a52fd33 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -1971,12 +1971,6 @@ def get_correct_attn_implementation(self, requested_attention: str | None, is_in
 
     def get_correct_experts_implementation(self, requested_experts: str | None) -> str:
         applicable_experts = "grouped_mm" if requested_experts is None else requested_experts
-        if applicable_experts not in ["eager", "grouped_mm", "batched_mm", "deepgemm", "sonicmoe"]:
-            message = (
-                f'Specified `experts_implementation="{applicable_experts}"` is not supported. The only possible arguments are '
-                '`experts_implementation="eager"`, `"experts_implementation=grouped_mm"`, `"experts_implementation=batched_mm"`, '
-                '`experts_implementation=deepgemm`, `"experts_implementation=sonicmoe"`'
-            )
         base_experts_fns = ["eager"] + list(set(ALL_EXPERTS_FUNCTIONS.keys()) | set(ALL_FP8_EXPERTS_FUNCTIONS.keys()))
         valid_experts_str_list = [f'`experts_implementation="{fn}"`' for fn in base_experts_fns]
         valid_experts_str_list[-1] = "and " + valid_experts_str_list[-1]

From a4c893451ba829419e0fe5d5c362187f9e4e27f5 Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <ilyas.moutawwakil@gmail.com>
Date: Thu, 23 Apr 2026 11:54:41 +0000
Subject: [PATCH 08/13] simpler naming

---
 src/transformers/integrations/moe.py                | 6 +++---
 src/transformers/integrations/sonicmoe.py           | 8 +++-----
 src/transformers/models/gpt_oss/modeling_gpt_oss.py | 2 +-
 src/transformers/models/gpt_oss/modular_gpt_oss.py  | 2 +-
 4 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py
index 650030c04ef0..c8a8e87f3621 100644
--- a/src/transformers/integrations/moe.py
+++ b/src/transformers/integrations/moe.py
@@ -501,7 +501,7 @@ def use_experts_implementation(
     experts_class: type[torch.nn.Module] | None = None,
     *,
     experts_interface: ExpertsInterface = ALL_EXPERTS_FUNCTIONS,
-    is_concatenated_layout: bool = True,
+    is_concatenated: bool = True,
     is_transposed: bool = False,
     has_bias: bool = False,
     has_gate: bool = True,
@@ -513,7 +513,7 @@ def use_experts_implementation(
             The experts class to modify. If not provided, returns a decorator that can be applied to the class.
         experts_interface (`ExpertsInterface`, *optional*, defaults to `ALL_EXPERTS_FUNCTIONS`):
             The experts interface to use for dispatching the forward method.
-        is_concatenated_layout (`bool`, *optional*, defaults to `True`):
+        is_concatenated (`bool`, *optional*, defaults to `True`):
             Whether the expert weights are stored in concatenated layout [gate;up]
             or interleaved layout [gate0, up0, gate1, up1, ...].
         is_transposed (`bool`, *optional*, defaults to `False`):
@@ -539,7 +539,7 @@ def __init__(self, config, *args, **kwargs):
             self.has_gate = has_gate
             self.has_bias = has_bias
             self.is_transposed = is_transposed
-            self.is_concatenated_layout = is_concatenated_layout
+            self.is_concatenated = is_concatenated
 
         @wraps(original_forward)
         def forward(self, *args, **kwargs):
diff --git a/src/transformers/integrations/sonicmoe.py b/src/transformers/integrations/sonicmoe.py
index 7250ccd31640..04bed9b289e2 100644
--- a/src/transformers/integrations/sonicmoe.py
+++ b/src/transformers/integrations/sonicmoe.py
@@ -88,8 +88,6 @@ def sonicmoe_experts_forward(
     device = hidden_states.device
     num_top_k = top_k_index.size(-1)
     num_tokens = hidden_states.size(0)
-    grad_enabled = torch.is_grad_enabled()
-    stream_id = torch.cuda.current_stream(device).cuda_stream  # this breaks torch.compile on some versions
 
     # Flatten — token_indices must be int32, sorted ascending (required by sonic-moe)
     token_idx = torch.arange(num_tokens, device=device).unsqueeze(1).expand(-1, num_top_k).reshape(-1).int()
@@ -116,10 +114,10 @@ def sonicmoe_experts_forward(
         w2,
         b2,
         E=self.num_experts,
-        stream_id=stream_id,
         activation_type=activation_type,
-        is_inference_mode_enabled=not grad_enabled,
-        concat_layout=self.is_concatenated_layout,
+        stream_id=torch.cuda.current_stream(device).cuda_stream,
+        is_inference_mode_enabled=not torch.is_grad_enabled(),
+        concat_layout=self.is_concatenated,
     )
 
     return output
diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py
index fa269770258f..55381a7e3c21 100644
--- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py
+++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py
@@ -65,7 +65,7 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
-@use_experts_implementation(is_concatenated_layout=False, is_transposed=True, has_bias=True)
+@use_experts_implementation(is_concatenated=False, is_transposed=True, has_bias=True)
 class GptOssExperts(nn.Module):
     def __init__(self, config):
         super().__init__()
diff --git a/src/transformers/models/gpt_oss/modular_gpt_oss.py b/src/transformers/models/gpt_oss/modular_gpt_oss.py
index a151ecbf94c9..3354acef2196 100644
--- a/src/transformers/models/gpt_oss/modular_gpt_oss.py
+++ b/src/transformers/models/gpt_oss/modular_gpt_oss.py
@@ -62,7 +62,7 @@ def forward(self, hidden_states):
         return (self.weight * hidden_states).to(input_dtype)  # main diff with Llama
 
 
-@use_experts_implementation(is_concatenated_layout=False, is_transposed=True, has_bias=True)
+@use_experts_implementation(is_concatenated=False, is_transposed=True, has_bias=True)
 class GptOssExperts(nn.Module):
     def __init__(self, config):
         super().__init__()

From 09910a1315d785eadd5b655b9dc02a5f968e3951 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Thu, 23 Apr 2026 14:00:54 +0200
Subject: [PATCH 09/13] style

---
 .../openai_privacy_filter/modeling_openai_privacy_filter.py     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py
index 4082a3ee9bf4..dbd1f1dda350 100644
--- a/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py
+++ b/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py
@@ -258,7 +258,7 @@ def forward(
         return attn_output, attn_weights
 
 
-@use_experts_implementation(is_transposed=True, has_bias=True)
+@use_experts_implementation(is_concatenated=False, is_transposed=True, has_bias=True)
 class OpenAIPrivacyFilterExperts(nn.Module):
     def __init__(self, config):
         super().__init__()

From 76abdc2722f211384b19c83799cf6e38ce65b300 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Thu, 23 Apr 2026 14:27:36 +0200
Subject: [PATCH 10/13] add sonicmoe test

---
 tests/test_modeling_common.py | 94 ++++++++++++++++-------------------
 1 file changed, 42 insertions(+), 52 deletions(-)

diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
index b909212b62cd..3afb2df0c0d5 100644
--- a/tests/test_modeling_common.py
+++ b/tests/test_modeling_common.py
@@ -51,7 +51,11 @@
     is_deepspeed_zero3_enabled,
     unset_hf_deepspeed_config,
 )
-from transformers.integrations.moe import batched_mm_experts_forward, grouped_mm_experts_forward
+from transformers.integrations.moe import (
+    batched_mm_experts_forward,
+    grouped_mm_experts_forward,
+    sonicmoe_experts_forward,
+)
 from transformers.modeling_layers import GradientCheckpointingLayer
 from transformers.modeling_utils import FLASH_ATTN_KERNEL_FALLBACK, _get_tied_weight_keys
 from transformers.models.auto import get_values
@@ -110,6 +114,7 @@
     GENERATION_CONFIG_NAME,
     SAFE_WEIGHTS_NAME,
     ModelOutput,
+    is_kernels_available,
     is_torch_bf16_available_on_device,
     is_torch_fp16_available_on_device,
 )
@@ -576,59 +581,44 @@ def _test_eager_matches_batched_and_grouped_inference(self, name, dtype):
             model.save_pretrained(tmpdirname)
             model = model_class.from_pretrained(tmpdirname).eval().to(torch_device).to(dtype)
 
-        with torch.no_grad():
-            inputs_dict = {k: v.to(dtype) if torch.is_floating_point(v) else v for k, v in inputs_dict.items()}
-            prepared_inputs = self._prepare_for_class(inputs_dict, model_class)
-
-            mock_batched_mm_forward = Mock(wraps=batched_mm_experts_forward)
-            mock_grouped_mm_forward = Mock(wraps=grouped_mm_experts_forward)
-            with (
-                # This is needed because we call the functions through the interface's global mapping
-                patch.dict(
-                    "transformers.integrations.moe.ALL_EXPERTS_FUNCTIONS._global_mapping",
-                    {"batched_mm": mock_batched_mm_forward, "grouped_mm": mock_grouped_mm_forward},
-                ),
-            ):
-                model.set_experts_implementation("eager")
-                self.assertEqual(model.config._experts_implementation, "eager")
-                outputs_eager = model(**prepared_inputs)
-                mock_batched_mm_forward.assert_not_called()
-                mock_grouped_mm_forward.assert_not_called()
+        inputs_dict = {k: v.to(dtype) if torch.is_floating_point(v) else v for k, v in inputs_dict.items()}
+        prepared_inputs = self._prepare_for_class(inputs_dict, model_class)
 
-                mock_batched_mm_forward.reset_mock()
-                mock_grouped_mm_forward.reset_mock()
+        implementations = ["eager", "batched_mm", "grouped_mm"]
+        mocks = {
+            "batched_mm": Mock(wraps=batched_mm_experts_forward),
+            "grouped_mm": Mock(wraps=grouped_mm_experts_forward),
+        }
 
-                model.set_experts_implementation("batched_mm")
-                self.assertEqual(model.config._experts_implementation, "batched_mm")
-                outputs_batched_mm = model(**prepared_inputs)
-                mock_grouped_mm_forward.assert_not_called()
-                mock_batched_mm_forward.assert_called()
-
-                mock_batched_mm_forward.reset_mock()
-                mock_grouped_mm_forward.reset_mock()
-
-                model.set_experts_implementation("grouped_mm")
-                self.assertEqual(model.config._experts_implementation, "grouped_mm")
-                outputs_grouped_mm = model(**prepared_inputs)
-                mock_batched_mm_forward.assert_not_called()
-                mock_grouped_mm_forward.assert_called()
-
-                mock_batched_mm_forward.reset_mock()
-                mock_grouped_mm_forward.reset_mock()
-
-        # extract output tensors for comparison
-        outputs_eager = _get_output_tensors(outputs_eager)
-        outputs_batched_mm = _get_output_tensors(outputs_batched_mm)
-        outputs_grouped_mm = _get_output_tensors(outputs_grouped_mm)
-
-        # make sure we have collected some tensors from the outputs
-        self.assertTrue(outputs_eager, "No outputs from eager implementation")
-        self.assertTrue(outputs_batched_mm, "No outputs from batched_mm implementation")
-        self.assertTrue(outputs_grouped_mm, "No outputs from grouped_mm implementation")
-
-        # make sure all implementations give numerically close outputs
-        torch.testing.assert_close(outputs_eager, outputs_batched_mm, rtol=1e-4, atol=1e-4)
-        torch.testing.assert_close(outputs_eager, outputs_grouped_mm, rtol=1e-4, atol=1e-4)
+        if is_kernels_available() and torch.cuda.is_available() and torch.cuda.get_device_capability() >= (9, 0):
+            # we also need nvidia-cutlass-dsl and apache-tvm-ffi
+            mocks["sonicmoe"] = Mock(wraps=sonicmoe_experts_forward)
+            implementations.append("sonicmoe")
+
+        outputs = {}
+        # This is needed because we call the functions through the interface's global mapping
+        with patch.dict("transformers.integrations.moe.ALL_EXPERTS_FUNCTIONS._global_mapping", mocks):
+            for impl in implementations:
+                model.set_experts_implementation(impl)
+                self.assertEqual(model.config._experts_implementation, impl)
+
+                with torch.no_grad():
+                    outputs[impl] = _get_output_tensors(model(**prepared_inputs))
+
+                self.assertTrue(outputs[impl], f"No outputs from {impl} implementation")
+
+                for name, mock in mocks.items():
+                    if name == impl:
+                        mock.assert_called()
+                    else:
+                        mock.assert_not_called()
+
+                    mock.reset_mock()
+
+        # all non-eager implementations must numerically match eager
+        eager_outputs = outputs.pop("eager")
+        for impl, impl_outputs in outputs.items():
+            torch.testing.assert_close(eager_outputs, impl_outputs, rtol=1e-4, atol=1e-4)
 
 
 def _config_zero_init(config):

From 2007f78a2af8e4320a49b8dd18b8d349b7b2b7e1 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Thu, 23 Apr 2026 14:41:26 +0200
Subject: [PATCH 11/13] skip fp32 on sonic

---
 tests/test_modeling_common.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
index 3afb2df0c0d5..bc8f65891445 100644
--- a/tests/test_modeling_common.py
+++ b/tests/test_modeling_common.py
@@ -590,7 +590,12 @@ def _test_eager_matches_batched_and_grouped_inference(self, name, dtype):
             "grouped_mm": Mock(wraps=grouped_mm_experts_forward),
         }
 
-        if is_kernels_available() and torch.cuda.is_available() and torch.cuda.get_device_capability() >= (9, 0):
+        if (
+            dtype != torch.float32
+            and is_kernels_available()
+            and torch.cuda.is_available()
+            and torch.cuda.get_device_capability() >= (9, 0)
+        ):
             # we also need nvidia-cutlass-dsl and apache-tvm-ffi
             mocks["sonicmoe"] = Mock(wraps=sonicmoe_experts_forward)
             implementations.append("sonicmoe")

From 1669103797e60197c5011daa07fcbfc49dacdf48 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Thu, 23 Apr 2026 14:55:23 +0200
Subject: [PATCH 12/13] add transposed support

---
 src/transformers/integrations/sonicmoe.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/transformers/integrations/sonicmoe.py b/src/transformers/integrations/sonicmoe.py
index 04bed9b289e2..e322bb4bc061 100644
--- a/src/transformers/integrations/sonicmoe.py
+++ b/src/transformers/integrations/sonicmoe.py
@@ -15,7 +15,7 @@
 """SonicMoE integration: fused MoE using CuteDSL kernels from `kernels-community/sonic-moe`.
 
 Provides `sonicmoe_experts_forward` registered as "sonicmoe" in the ExpertsInterface.
-Requirements: CUDA, `kernels`, `nvidia-cutlass-dsl`, has_gate=True, is_transposed=False.
+Requirements: CUDA, `kernels`, `nvidia-cutlass-dsl`, has_gate=True.
 """
 
 import functools
@@ -76,8 +76,6 @@ def sonicmoe_experts_forward(
     top_k_index: torch.Tensor,
     top_k_weights: torch.Tensor,
 ) -> torch.Tensor:
-    if self.is_transposed:
-        raise ValueError("sonicmoe requires non-transposed weights (is_transposed=False)")
     if not self.has_gate:
         raise ValueError("sonicmoe requires gated experts (has_gate=True)")
     if hidden_states.device.type != "cuda":
@@ -98,9 +96,12 @@ def sonicmoe_experts_forward(
     act_name = getattr(self.config, "hidden_act", "silu").lower()
     activation_type = getattr(ActivationType, ACT_MAP.get(act_name, "swiglu").upper(), ActivationType.SWIGLU)
 
-    # Permute weights as expected by sonic-moe (E=num_experts, H=hidden_size, I=intermediate_size)
-    w1 = self.gate_up_proj.permute(1, 2, 0)  # (2*I, H, E)
-    w2 = self.down_proj.permute(1, 2, 0)  # (I, H, E)
+    # Permute weights as expected by sonic-moe (E=num_experts, H=hidden_size, I=intermediate_size).
+    # Non-transposed: gate_up_proj is (E, 2*I, H), down_proj is (E, H, I) -> permute(1, 2, 0).
+    # Transposed: gate_up_proj is (E, H, 2*I), down_proj is (E, I, H) -> permute(2, 1, 0).
+    perm = (2, 1, 0) if self.is_transposed else (1, 2, 0)
+    w1 = self.gate_up_proj.permute(*perm)  # (2*I, H, E)
+    w2 = self.down_proj.permute(*perm)  # (I, H, E)
     b1 = self.gate_up_proj_bias if self.has_bias else None
     b2 = self.down_proj_bias if self.has_bias else None
 

From 30b915b572ede7941fd8db48ce08af2d87267f1c Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Thu, 23 Apr 2026 15:08:15 +0200
Subject: [PATCH 13/13] fix

---
 .../openai_privacy_filter/modeling_openai_privacy_filter.py     | 2 +-
 .../openai_privacy_filter/modular_openai_privacy_filter.py      | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py
index dbd1f1dda350..4082a3ee9bf4 100644
--- a/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py
+++ b/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py
@@ -258,7 +258,7 @@ def forward(
         return attn_output, attn_weights
 
 
-@use_experts_implementation(is_concatenated=False, is_transposed=True, has_bias=True)
+@use_experts_implementation(is_transposed=True, has_bias=True)
 class OpenAIPrivacyFilterExperts(nn.Module):
     def __init__(self, config):
         super().__init__()
diff --git a/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py
index fc77aafbdcf5..422235d9da91 100644
--- a/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py
+++ b/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py
@@ -21,6 +21,7 @@
 from torch.nn import functional as F
 
 from ...configuration_utils import PreTrainedConfig
+from ...integrations import use_experts_implementation
 from ...masking_utils import create_bidirectional_sliding_window_mask
 from ...modeling_layers import GenericForTokenClassification
 from ...modeling_outputs import BaseModelOutput
@@ -213,6 +214,7 @@ def forward(
         return attn_output, attn_weights
 
 
+@use_experts_implementation(is_transposed=True, has_bias=True)
 class OpenAIPrivacyFilterExperts(GptOssExperts):
     def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor:
         # Concatenated layout instead of interleaving