Add AFMoE model implementation (#13216)

docs/supported_models/generative_models.md

@@ -64,3 +64,4 @@ in the GitHub search bar.
 | **NVIDIA Nemotron Nano 2.0** | `nvidia/NVIDIA-Nemotron-Nano-9B-v2` | The [NVIDIA Nemotron](https://www.nvidia.com/en-us/ai-data-science/foundation-models/nemotron/) family of multimodal models provides state-of-the-art reasoning models specifically designed for enterprise-ready AI agents. `Nemotron-Nano-9B-v2` is a hybrid Mamba-Transformer language model designed to increase throughput for reasoning workloads while achieving state-of-the-art accuracy compared to similarly-sized models. |
 | **StarCoder2** (3B-15B)               | `bigcode/starcoder2-7b`                         | StarCoder2 is a family of open large language models (LLMs) specialized for code generation and understanding. It is the successor to StarCoder, jointly developed by the BigCode project (a collaboration between Hugging Face, ServiceNow Research, and other contributors). |
 | **Jet-Nemotron** | `jet-ai/Jet-Nemotron-2B` | Jet-Nemotron is a new family of hybrid-architecture language models that surpass state-of-the-art open-source full-attention language models, while achieving significant efficiency gains. |
+| **Trinity** (Nano, Mini) | `arcee-ai/Trinity-Mini` | Arcee's foundational MoE Trinity family of models, open weights under Apache 2.0. |

python/sglang/srt/configs/__init__.py

@@ -1,3 +1,4 @@
+from sglang.srt.configs.afmoe import AfmoeConfig
 from sglang.srt.configs.chatglm import ChatGLMConfig
 from sglang.srt.configs.dbrx import DbrxConfig
 from sglang.srt.configs.deepseekvl2 import DeepseekVL2Config
@@ -23,6 +24,7 @@ from sglang.srt.configs.step3_vl import (
 )
 
 __all__ = [
+    "AfmoeConfig",
     "ExaoneConfig",
     "ChatGLMConfig",
     "DbrxConfig",

python/sglang/srt/configs/afmoe.py

@@ -0,0 +1,102 @@
+from typing import List, Optional
+
+from transformers import PretrainedConfig
+
+
+class AfmoeConfig(PretrainedConfig):
+    model_type = "afmoe"
+
+    def __init__(
+        self,
+        vocab_size: int = 32000,
+        hidden_size: int = 4096,
+        intermediate_size: int = 11008,
+        moe_intermediate_size: int = 256,
+        num_hidden_layers: int = 32,
+        num_attention_heads: int = 32,
+        num_key_value_heads: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 131072,
+        initializer_range: float = 0.02,
+        rms_norm_eps: float = 1e-5,
+        use_cache: bool = True,
+        pad_token_id: Optional[int] = None,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        tie_word_embeddings: bool = False,
+        rope_theta: float = 10000.0,
+        rope_scaling: Optional[dict] = None,
+        attention_bias: bool = False,
+        attention_dropout: float = 0.0,
+        # MoE parameters
+        num_experts: Optional[int] = None,
+        num_experts_per_tok: Optional[int] = None,
+        num_shared_experts: int = 0,
+        num_dense_layers: int = 0,
+        # Routing parameters
+        score_func: str = "sigmoid",
+        route_norm: bool = True,
+        route_scale: float = 1.0,
+        n_group: int = 1,
+        topk_group: int = 1,
+        # Attention parameters
+        sliding_window: Optional[int] = None,
+        layer_types: Optional[List[str]] = None,
+        global_attn_every_n_layers: int = 4,
+        # muP scaling
+        mup_enabled: bool = False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.moe_intermediate_size = moe_intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = (
+            head_dim if head_dim is not None else hidden_size // num_attention_heads
+        )
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+
+        # MoE parameters
+        self.num_experts = num_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.num_shared_experts = num_shared_experts
+        self.num_dense_layers = num_dense_layers
+
+        # Routing parameters
+        self.score_func = score_func
+        self.route_norm = route_norm
+        self.route_scale = route_scale
+        self.n_group = n_group
+        self.topk_group = topk_group
+
+        # Attention parameters
+        self.sliding_window = sliding_window
+        self.layer_types = layer_types
+        self.global_attn_every_n_layers = global_attn_every_n_layers
+
+        # muP scaling
+        self.mup_enabled = mup_enabled
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

python/sglang/srt/models/afmoe.py

@@ -0,0 +1,633 @@
+# Copyright 2023-2024 SGLang Team
+# 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.
+# ==============================================================================
+"""Inference-only AfMoE model compatible with HuggingFace weights.
+
+AfMoE is a Mixture-of-Experts model with:
+- Gated attention with sigmoid gating
+- Q/K normalization with RMSNorm
+- Dual normalization (pre/post for both attention and MLP)
+- Sliding window attention for local layers
+- muP (maximal update parameterization) scaling support
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import Iterable, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import PretrainedConfig
+
+from sglang.srt.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
+)
+from sglang.srt.layers.activation import SiluAndMul
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.srt.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from sglang.srt.layers.logits_processor import LogitsProcessor
+from sglang.srt.layers.moe.fused_moe_triton import fused_moe
+from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
+from sglang.srt.layers.moe.topk import TopK
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.rotary_embedding import get_rope
+from sglang.srt.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.utils import add_prefix
+
+
+def get_attention_sliding_window_size(config: PretrainedConfig) -> Optional[int]:
+    sliding_window = getattr(config, "sliding_window", None)
+    if sliding_window is None:
+        return None
+    if sliding_window <= 0:
+        return None
+    # Align with other local attention implementations (see gpt_oss).
+    return sliding_window - 1
+
+
+class AfmoeMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: Optional[QuantizationConfig] = None,
+        reduce_results: bool = True,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            hidden_size,
+            [intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("gate_up_proj", prefix),
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
+            bias=False,
+            reduce_results=reduce_results,
+            quant_config=quant_config,
+            prefix=add_prefix("down_proj", prefix),
+        )
+        if hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
+            )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class AfmoeMoE(nn.Module):
+
+    @staticmethod
+    def _custom_routing_function(
+        hidden_states: torch.Tensor,
+        gating_output: torch.Tensor,
+        topk: int,
+        renormalize: bool,
+        *,
+        score_func: str,
+        expert_bias: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        logits = gating_output.to(torch.float32)
+        if score_func == "sigmoid":
+            scores = torch.sigmoid(logits)
+            if expert_bias is not None:
+                bias = expert_bias.to(scores.device, dtype=scores.dtype)
+                scores_for_choice = scores + bias
+                topk_ids = torch.topk(scores_for_choice, k=topk, dim=-1)[1]
+                topk_weights = scores.gather(dim=-1, index=topk_ids)
+            else:
+                topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1)
+        else:
+            if expert_bias is not None:
+                logits = logits + expert_bias.to(logits.device, dtype=logits.dtype)
+            probs = F.softmax(logits, dim=-1)
+            topk_weights, topk_ids = torch.topk(probs, k=topk, dim=-1)
+
+        if renormalize:
+            denom = topk_weights.sum(dim=-1, keepdim=True).clamp(min=1e-20)
+            topk_weights = topk_weights / denom
+
+        return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.config = config
+        self.rank = get_tensor_model_parallel_rank()
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+        self.n_routed_experts = getattr(config, "num_experts", None)
+        if self.n_routed_experts is None:
+            raise ValueError("AfmoeConfig must define `num_experts`.")
+        self.top_k = config.num_experts_per_tok
+        if self.tp_size > self.n_routed_experts:
+            raise ValueError(
+                f"Tensor parallel size {self.tp_size} is greater than "
+                f"the number of experts {self.n_routed_experts}."
+            )
+
+        self.score_func = getattr(config, "score_func", "softmax")
+        self.route_norm = getattr(config, "route_norm", True)
+        self.route_scale = float(getattr(config, "route_scale", 1.0))
+        self.n_group = getattr(config, "n_group", 1)
+        self.topk_group = getattr(config, "topk_group", 1)
+        self.use_grouped_topk = self.n_group is not None and self.n_group > 1
+        self.num_shared_experts = getattr(config, "num_shared_experts", 0)
+
+        self.gate = ReplicatedLinear(
+            config.hidden_size,
+            self.n_routed_experts,
+            bias=False,
+            quant_config=None,
+            prefix=add_prefix("gate", prefix),
+        )
+
+        self.expert_bias = nn.Parameter(
+            torch.zeros(self.n_routed_experts, dtype=torch.float32),
+            requires_grad=False,
+        )
+
+        self.experts = nn.ModuleList(
+            [
+                AfmoeMLP(
+                    hidden_size=config.hidden_size,
+                    intermediate_size=config.moe_intermediate_size,
+                    hidden_act=config.hidden_act,
+                    quant_config=quant_config,
+                    reduce_results=False,
+                    prefix=add_prefix(f"experts.{idx}", prefix),
+                )
+                for idx in range(self.n_routed_experts)
+            ]
+        )
+        self.pack_params()
+
+        if self.num_shared_experts:
+            intermediate_size = config.moe_intermediate_size * self.num_shared_experts
+            self.shared_experts = AfmoeMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                reduce_results=False,
+                prefix=add_prefix("shared_experts", prefix),
+            )
+        else:
+            self.shared_experts = None
+
+        custom_routing_fn = None
+        correction_bias = None
+        if self.use_grouped_topk:
+            correction_bias = self.expert_bias
+        elif self.score_func == "sigmoid":
+            custom_routing_fn = functools.partial(
+                AfmoeMoE._custom_routing_function,
+                score_func=self.score_func,
+                expert_bias=self.expert_bias,
+            )
+
+        renormalize = self.route_norm if self.score_func == "sigmoid" else False
+        self.topk = TopK(
+            top_k=self.top_k,
+            renormalize=renormalize,
+            use_grouped_topk=self.use_grouped_topk,
+            num_expert_group=self.n_group if self.use_grouped_topk else None,
+            topk_group=self.topk_group if self.use_grouped_topk else None,
+            custom_routing_function=custom_routing_fn,
+            correction_bias=correction_bias,
+            routed_scaling_factor=self.route_scale,
+        )
+
+    def pack_params(self) -> None:
+        w1: list[torch.Tensor] = []
+        w2: list[torch.Tensor] = []
+        for expert in self.experts:
+            w1.append(expert.gate_up_proj.weight)
+            w2.append(expert.down_proj.weight)
+        self.w1 = torch._utils._flatten_dense_tensors(w1)
+        w1s = torch._utils._unflatten_dense_tensors(self.w1, w1)
+        for data, param in zip(w1s, w1):
+            param.data = data
+        self.w1 = self.w1.view(len(w1), *w1s[0].shape)
+
+        self.w2 = torch._utils._flatten_dense_tensors(w2)
+        w2s = torch._utils._unflatten_dense_tensors(self.w2, w2)
+        for data, param in zip(w2s, w2):
+            param.data = data
+        self.w2 = self.w2.view(len(w2), *w2s[0].shape)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        shared_output = None
+        if self.shared_experts is not None:
+            shared_output = self.shared_experts(hidden_states)
+
+        router_logits, _ = self.gate(hidden_states)
+        topk_output = self.topk(hidden_states, router_logits)
+        final_hidden_states = fused_moe.fused_moe(
+            hidden_states,
+            w1=self.w1,
+            w2=self.w2,
+            topk_output=topk_output,
+            moe_runner_config=MoeRunnerConfig(
+                inplace=True,
+                routed_scaling_factor=self.route_scale,
+            ),
+        )
+
+        if shared_output is not None:
+            final_hidden_states = final_hidden_states + shared_output
+
+        final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
+        return final_hidden_states.view(num_tokens, hidden_dim)
+
+
+class AfmoeAttention(nn.Module):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        layer_id: int = 0,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+
+        self.head_dim = getattr(config, "head_dim", hidden_size // self.total_num_heads)
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+
+        rope_theta = getattr(config, "rope_theta", 10000)
+        rope_scaling = getattr(config, "rope_scaling", None)
+        partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
+        self.rotary_dim = int(self.head_dim * partial_rotary_factor)
+        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+
+        layer_types = getattr(config, "layer_types", None)
+        self.is_local_attention = (
+            layer_types is not None and layer_types[layer_id] == "sliding_attention"
+        )
+        sliding_window = (
+            get_attention_sliding_window_size(config) if self.is_local_attention else -1
+        )
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("qkv_proj", prefix),
+        )
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("o_proj", prefix),
+        )
+        self.gate_proj = ColumnParallelLinear(
+            hidden_size,
+            self.total_num_heads * self.head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("gate_proj", prefix),
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.rotary_dim,
+            max_position=max_position_embeddings,
+            base=rope_theta,
+            rope_scaling=rope_scaling,
+            is_neox_style=True,
+        )
+        self.attn = RadixAttention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            layer_id=layer_id,
+            sliding_window_size=sliding_window,
+            quant_config=quant_config,
+            prefix=add_prefix("attn", prefix),
+        )
+
+        eps = getattr(config, "rms_norm_eps", 1e-5)
+        self.q_norm = RMSNorm(self.head_dim, eps=eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=eps)
+        self.sliding_window = sliding_window
+
+    def _apply_qk_norm(
+        self, q: torch.Tensor, k: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        q_heads = self.q_norm(q.reshape(-1, self.head_dim))
+        k_heads = self.k_norm(k.reshape(-1, self.head_dim))
+        q = q_heads.view(q.shape)
+        k = k_heads.view(k.shape)
+        return q, k
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self._apply_qk_norm(q, k)
+
+        if self.is_local_attention:
+            q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v, forward_batch)
+
+        gate_vals, _ = self.gate_proj(hidden_states)
+        attn_output = attn_output * torch.sigmoid(gate_vals)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class AfmoeDecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        layer_id: int,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.layer_id = layer_id
+
+        self.self_attn = AfmoeAttention(
+            config=config,
+            hidden_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            layer_id=layer_id,
+            quant_config=quant_config,
+            prefix=add_prefix("self_attn", prefix),
+        )
+
+        use_moe = False
+        if hasattr(config, "num_dense_layers"):
+            use_moe = layer_id >= config.num_dense_layers
+        elif (
+            getattr(config, "num_experts", None) is not None
+            and hasattr(config, "first_k_dense_replace")
+            and hasattr(config, "moe_layer_freq")
+        ):
+            base = config.first_k_dense_replace
+            freq = config.moe_layer_freq
+            use_moe = layer_id >= base and (layer_id - base) % freq == 0
+
+        if use_moe:
+            self.mlp = AfmoeMoE(
+                config=config,
+                quant_config=quant_config,
+                prefix=add_prefix("mlp", prefix),
+            )
+        else:
+            self.mlp = AfmoeMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=add_prefix("mlp", prefix),
+            )
+
+        eps = getattr(config, "rms_norm_eps", 1e-5)
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=eps)
+        self.pre_mlp_layernorm = RMSNorm(config.hidden_size, eps=eps)
+        self.post_mlp_layernorm = RMSNorm(config.hidden_size, eps=eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+        attn_residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(positions, hidden_states, forward_batch)
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = attn_residual + hidden_states
+
+        mlp_residual = hidden_states
+        hidden_states = self.pre_mlp_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.post_mlp_layernorm(hidden_states)
+        hidden_states = mlp_residual + hidden_states
+
+        return hidden_states
+
+
+class AfmoeModel(nn.Module):
+
+    fall_back_to_pt_during_load = False
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+        )
+        self.layers = nn.ModuleList(
+            [
+                AfmoeDecoderLayer(
+                    config,
+                    layer_id,
+                    quant_config=quant_config,
+                    prefix=add_prefix(f"layers.{layer_id}", prefix),
+                )
+                for layer_id in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        input_embeds: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if input_embeds is None:
+            hidden_states = self.embed_tokens(input_ids)
+        else:
+            hidden_states = input_embeds
+
+        if getattr(self.config, "mup_enabled", False):
+            hidden_states = hidden_states * (self.config.hidden_size**0.5)
+
+        for layer in self.layers:
+            hidden_states = layer(positions, hidden_states, forward_batch)
+        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed_tokens
+
+
+class AfmoeForCausalLM(nn.Module):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.quant_config = quant_config
+        self.model = AfmoeModel(
+            config, quant_config, prefix=add_prefix("model", prefix)
+        )
+        self.lm_head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=add_prefix("lm_head", prefix),
+        )
+        self.logits_processor = LogitsProcessor(config)
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.embed_tokens
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        input_embeds: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        hidden_states = self.model(input_ids, positions, forward_batch, input_embeds)
+        return self.logits_processor(
+            input_ids, hidden_states, self.lm_head, forward_batch
+        )
+
+    def get_attention_sliding_window_size(self) -> Optional[int]:
+        return get_attention_sliding_window_size(self.config)
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None:
+        stacked_params_mapping = [
+            # (param_name, weight_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        for name, loaded_weight in weights:
+            # Skip rotary embedding inverse frequencies
+            if "rotary_emb.inv_freq" in name:
+                continue
+
+            # Remap router gate weights: HF uses .mlp.router.gate., SGLang uses .mlp.gate.
+            if ".mlp.router.gate." in name:
+                name = name.replace(".mlp.router.gate.", ".mlp.gate.")
+
+            # Handle stacked params (qkv_proj, gate_up_proj)
+            handled = False
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                # Skip gate_proj/up_proj stacking for self_attn (attention uses separate gate_proj)
+                if ".self_attn." in name and weight_name in {"gate_proj", "up_proj"}:
+                    continue
+
+                new_name = name.replace(weight_name, param_name)
+                # Skip if parameter doesn't exist (e.g., bias for layers without bias)
+                if new_name not in params_dict:
+                    handled = True
+                    break
+
+                param = params_dict[new_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight, shard_id)
+                handled = True
+                break
+
+            if handled:
+                continue
+
+            # Load remaining weights directly
+            if name in params_dict:
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+
+EntryClass = AfmoeForCausalLM

python/sglang/srt/utils/hf_transformers_utils.py

@@ -44,6 +44,7 @@ from transformers import (
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 
 from sglang.srt.configs import (
+    AfmoeConfig,
     ChatGLMConfig,
     DbrxConfig,
     DeepseekVL2Config,
@@ -71,6 +72,7 @@ from sglang.srt.utils import is_remote_url, logger, lru_cache_frozenset, mistral
 from sglang.srt.utils.patch_tokenizer import patch_tokenizer
 
 _CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [
+    AfmoeConfig,
     ChatGLMConfig,
     DbrxConfig,
     ExaoneConfig,