diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index b1e6c74ddf10..70a343424aa8 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -289,6 +289,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/moe.py b/src/transformers/integrations/moe.py index d17522d26daa..c8a8e87f3621 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -23,6 +23,7 @@ is_torch_less_or_equal, is_torchdynamo_compiling, ) +from .sonicmoe import sonicmoe_experts_forward if is_torch_available(): @@ -31,6 +32,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.""" @@ -458,6 +460,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, } @@ -498,6 +501,7 @@ def use_experts_implementation( experts_class: type[torch.nn.Module] | None = None, *, experts_interface: ExpertsInterface = ALL_EXPERTS_FUNCTIONS, + is_concatenated: bool = True, is_transposed: bool = False, has_bias: bool = False, has_gate: bool = True, @@ -509,10 +513,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 (`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. @@ -529,6 +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 = is_concatenated @wraps(original_forward) def forward(self, *args, **kwargs): diff --git a/src/transformers/integrations/sonicmoe.py b/src/transformers/integrations/sonicmoe.py new file mode 100644 index 000000000000..e322bb4bc061 --- /dev/null +++ b/src/transformers/integrations/sonicmoe.py @@ -0,0 +1,124 @@ +# 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. +""" + +import functools + +import torch + +from ..utils import logging +from .hub_kernels import lazy_load_kernel + + +logger = logging.get_logger(__name__) + +# Map activation function names from HF config to SonicMoE epilogue names +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, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + 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") + + ActivationType, moe_general_routing_inputs = _load_sonic_kernel() + + device = hidden_states.device + num_top_k = top_k_index.size(-1) + num_tokens = hidden_states.size(0) + + # 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() + router_scores = top_k_weights.reshape(-1).to(hidden_states.dtype) + 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 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 + + output, _ = moe_general_routing_inputs( + hidden_states, + router_scores, + token_idx, + expert_ids, + w1, + b1, + w2, + b2, + E=self.num_experts, + activation_type=activation_type, + 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 00f9ac601b0e..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_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 f7c89cab08e5..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_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/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 diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index b909212b62cd..bc8f65891445 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,49 @@ 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 ( + 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") + + 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):