Model: Support IBM Granite (Dense/Mamba + MoE) (#18040)

python/sglang/srt/configs/__init__.py

@@ -7,6 +7,7 @@ from sglang.srt.configs.dots_ocr import DotsOCRConfig
 from sglang.srt.configs.dots_vlm import DotsVLMConfig
 from sglang.srt.configs.exaone import ExaoneConfig
 from sglang.srt.configs.falcon_h1 import FalconH1Config
+from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig
 from sglang.srt.configs.janus_pro import MultiModalityConfig
 from sglang.srt.configs.jet_nemotron import JetNemotronConfig
 from sglang.srt.configs.jet_vlm import JetVLMConfig
@@ -52,6 +53,7 @@ __all__ = [
     "DotsVLMConfig",
     "DotsOCRConfig",
     "FalconH1Config",
+    "GraniteMoeHybridConfig",
     "Lfm2Config",
     "Lfm2MoeConfig",
     "NemotronHConfig",

python/sglang/srt/configs/granitemoehybrid.py

@@ -0,0 +1,301 @@
+# coding=utf-8
+# Copyright 2025 IBM and the HuggingFace Inc. 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.
+"""GraniteMoeHybrid model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
+
+logger = logging.get_logger(__name__)
+
+MAMBA = "mamba"
+ATTENTION = "attention"
+
+
+class GraniteMoeHybridConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`GraniteMoeHybridModel`]. It is used to instantiate a
+    GraniteMoeHybrid model according to the specified arguments, defining the model architecture. The GraniteMoeHybrid is a
+    hybrid architecture combining Mamba2 layers with attention layers, developed by IBM.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 100352):
+            Vocabulary size of the GraniteMoeHybrid model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`GraniteMoeHybridModel`]
+        tie_word_embeddings (`bool`, *optional*, defaults to `True`):
+            Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the
+            model has a output word embedding layer.
+        hidden_size (`int`, *optional*, defaults to 2048):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 8192):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 40):
+            Number of hidden layers in the model.
+        layer_types (`list[str]`, *optional*):
+            List of layer types for each layer. Each element should be either "mamba" or "attention".
+            If not provided, defaults to alternating pattern based on num_hidden_layers.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        initializer_range (`float`, *optional*, defaults to 0.1):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon used by the rms normalization layers.
+        normalization_function (`str`, *optional*, defaults to `"rmsnorm"`):
+            The normalization function to use. Currently only "rmsnorm" is supported.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*, defaults to 100256):
+            The id of the padding token.
+        bos_token_id (`int`, *optional*, defaults to 100257):
+            The id of the "beginning-of-sequence" token.
+        eos_token_id (`int`, *optional*, defaults to 100257):
+            The id of the "end-of-sequence" token.
+        max_position_embeddings (`int`, *optional*, defaults to 131072):
+            Max cached sequence length for the model
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        attention_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use bias in attention layers.
+        position_embedding_type (`str`, *optional*, defaults to `"nope"`):
+            Type of position embedding. Can be "nope" (no position embedding) or "rope".
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The theta value used for the RoPE embeddings.
+        rope_scaling (`dict`, *optional*):
+            The scaling configuration for the RoPE embeddings. If `None`, no scaling is applied.
+        mamba_d_state (`int`, *optional*, defaults to 128):
+            The dimension of the mamba state space latents
+        mamba_d_conv (`int`, *optional*, defaults to 4):
+            The size of the mamba convolution kernel
+        mamba_expand (`int`, *optional*, defaults to 2):
+            Expanding factor (relative to hidden_size) used to determine the mamba intermediate size
+        mamba_d_head (`int`, *optional*, defaults to 64):
+            Head embedding dimension size for Mamba
+        mamba_n_heads (`int`, *optional*, defaults to 64):
+            The number of mamba heads
+        mamba_n_groups (`int`, *optional*, defaults to 1):
+            The number of the mamba groups
+        mamba_chunk_size (`int`, *optional*, defaults to 256):
+            The chunks in which to break the sequence when doing prefill/training
+        mamba_conv_bias (`bool`, *optional*, defaults to `True`):
+            Flag indicating whether or not to use bias in the convolution layer of the mamba mixer block.
+        mamba_proj_bias (`bool`, *optional*, defaults to `False`):
+            Flag indicating whether or not to use bias in the input and output projections of the mamba mixer block
+        embedding_multiplier (`float`, *optional*, defaults to 12.0):
+            The multiplier for the embedding layer. This is used to scale the output of the embedding layer.
+        logits_scaling (`float`, *optional*, defaults to 8.0):
+            The scaling factor for the logits.
+        attention_multiplier (`float`, *optional*, defaults to 0.015625):
+            The multiplier for the attention layers.
+        residual_multiplier (`float`, *optional*, defaults to 0.22):
+            The multiplier for residual connections.
+        num_local_experts (`int`, *optional*, defaults to 0):
+            Number of local experts in MoE layers.
+        num_experts_per_tok (`int`, *optional*, defaults to 0):
+            Number of experts to use per token in MoE layers.
+        shared_intermediate_size (`int`, *optional*, defaults to 8192):
+            Intermediate size for shared experts.
+        output_router_logits (`bool`, *optional*, defaults to `False`):
+            Whether to output router logits.
+        router_aux_loss_coef (`float`, *optional*, defaults to 0.01):
+            Auxiliary loss coefficient for the router.
+    """
+
+    model_type = "granitemoehybrid"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=100352,
+        tie_word_embeddings=True,
+        hidden_size=2048,
+        intermediate_size=8192,
+        num_hidden_layers=40,
+        layer_types=None,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        hidden_act="silu",
+        initializer_range=0.1,
+        rms_norm_eps=1e-5,
+        normalization_function="rmsnorm",
+        use_cache=True,
+        pad_token_id=100256,
+        bos_token_id=100257,
+        eos_token_id=100257,
+        max_position_embeddings=131072,
+        attention_dropout=0.0,
+        attention_bias=False,
+        position_embedding_type="nope",
+        rope_theta=10000.0,
+        rope_scaling=None,
+        mamba_d_state=128,
+        mamba_d_conv=4,
+        mamba_expand=2,
+        mamba_d_head=64,
+        mamba_n_heads=64,
+        mamba_n_groups=1,
+        mamba_chunk_size=256,
+        mamba_conv_bias=True,
+        mamba_proj_bias=False,
+        embedding_multiplier=12.0,
+        logits_scaling=8.0,
+        attention_multiplier=0.015625,
+        residual_multiplier=0.22,
+        num_local_experts=0,
+        num_experts_per_tok=0,
+        shared_intermediate_size=8192,
+        output_router_logits=False,
+        router_aux_loss_coef=0.01,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+
+        # Set layer types - if not provided, create default pattern
+        if layer_types is None:
+            # Default pattern: mamba layers with attention every 6th layer (roughly)
+            self.layer_types = []
+            for i in range(num_hidden_layers):
+                if (i + 1) % 6 == 0:
+                    self.layer_types.append(ATTENTION)
+                else:
+                    self.layer_types.append(MAMBA)
+        else:
+            self.layer_types = layer_types
+
+        # Validate layer_types
+        if len(self.layer_types) != self.num_hidden_layers:
+            raise ValueError(
+                f"layer_types must have length equal to num_hidden_layers ({num_hidden_layers}), "
+                f"but got {len(self.layer_types)}"
+            )
+
+        for layer_type in self.layer_types:
+            if layer_type not in [MAMBA, ATTENTION]:
+                raise ValueError(
+                    f"Each element in layer_types must be either '{MAMBA}' or '{ATTENTION}', "
+                    f"but got '{layer_type}'"
+                )
+
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.normalization_function = normalization_function
+
+        self.use_cache = use_cache
+        self.max_position_embeddings = max_position_embeddings
+        self.attention_dropout = attention_dropout
+        self.attention_bias = attention_bias
+
+        self.position_embedding_type = position_embedding_type
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+
+        # Mamba configuration
+        self.mamba_d_state = mamba_d_state
+        self.mamba_d_conv = mamba_d_conv
+        self.mamba_expand = mamba_expand
+        self.mamba_d_head = mamba_d_head
+        self.mamba_n_heads = mamba_n_heads
+        self.mamba_n_groups = mamba_n_groups
+        self.mamba_chunk_size = mamba_chunk_size
+        self.mamba_conv_bias = mamba_conv_bias
+        self.mamba_proj_bias = mamba_proj_bias
+
+        # Calculate mamba intermediate size
+        self.mamba_intermediate_size = mamba_expand * hidden_size
+
+        # Validate mamba configuration
+        if self.mamba_intermediate_size % mamba_n_heads != 0:
+            raise ValueError(
+                f"mamba_intermediate_size ({self.mamba_intermediate_size}) must be divisible by "
+                f"mamba_n_heads ({mamba_n_heads})"
+            )
+
+        if mamba_d_head * mamba_n_heads != self.mamba_intermediate_size:
+            raise ValueError(
+                f"mamba_d_head ({mamba_d_head}) * mamba_n_heads ({mamba_n_heads}) must equal "
+                f"mamba_intermediate_size ({self.mamba_intermediate_size})"
+            )
+
+        # Scaling factors
+        self.embedding_multiplier = embedding_multiplier
+        self.logits_scaling = logits_scaling
+        self.attention_multiplier = attention_multiplier
+        self.residual_multiplier = residual_multiplier
+
+        # MoE configuration
+        self.num_local_experts = num_local_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.shared_intermediate_size = shared_intermediate_size
+        self.output_router_logits = output_router_logits
+        self.router_aux_loss_coef = router_aux_loss_coef
+
+        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,
+        )
+
+    @property
+    def mamba_layer_ids(self):
+        """Returns the indices of layers that are Mamba layers."""
+        return [
+            i for i in range(self.num_hidden_layers) if self.layer_types[i] == MAMBA
+        ]
+
+    @property
+    def attention_layer_ids(self):
+        """Returns the indices of layers that are attention layers."""
+        return [
+            i for i in range(self.num_hidden_layers) if self.layer_types[i] == ATTENTION
+        ]
+
+    @property
+    def full_attention_layer_ids(self):
+        """Alias for attention_layer_ids for compatibility."""
+        return self.attention_layer_ids
+
+    @property
+    def mamba2_cache_params(self):
+        """Returns the Mamba2 cache parameters for this configuration."""
+        from sglang.srt.layers.dp_attention import get_attention_tp_size
+
+        shape = Mamba2StateShape.create(
+            tp_world_size=get_attention_tp_size(),
+            intermediate_size=self.mamba_intermediate_size,
+            n_groups=self.mamba_n_groups,
+            num_heads=self.mamba_n_heads,
+            head_dim=self.mamba_d_head,
+            state_size=self.mamba_d_state,
+            conv_kernel=self.mamba_d_conv,
+        )
+        return Mamba2CacheParams(shape=shape, layers=self.mamba_layer_ids)

python/sglang/srt/model_executor/model_runner.py

@@ -33,6 +33,7 @@ from torch import nn
 from sglang.srt.configs import (
     BailingHybridConfig,
     FalconH1Config,
+    GraniteMoeHybridConfig,
     JetNemotronConfig,
     JetVLMConfig,
     KimiLinearConfig,
@@ -1601,6 +1602,17 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             return config
         if isinstance(config, NemotronH_Nano_VL_V2_Config):
             return config.llm_config
+
+        if isinstance(config, GraniteMoeHybridConfig):
+            has_mamba = any(
+                layer_type == "mamba"
+                for layer_type in getattr(config, "layer_types", [])
+            )
+            if not has_mamba:
+                return None
+            else:
+                return config
+
         return None
 
     @property

python/sglang/srt/models/granitemoehybrid.py

@@ -0,0 +1,737 @@
+from typing import Iterable, Optional
+
+import torch
+from torch import nn
+from transformers.models.granitemoeshared import GraniteMoeSharedConfig
+
+from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig
+from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from sglang.srt.layers.activation import SiluAndMul
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+    HybridLinearAttnBackend,
+    Mamba2AttnBackend,
+)
+from sglang.srt.layers.attention.mamba.mamba import MambaMixer2
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.srt.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from sglang.srt.layers.logits_processor import LogitsProcessor
+from sglang.srt.layers.pooler import Pooler, PoolingType
+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.utils import PPMissingLayer
+from sglang.srt.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.models.transformers import maybe_prefix
+from sglang.srt.utils import make_layers
+
+from .granitemoe import GraniteMoeMoE
+
+
+# in vLLM this is in a separate file, but keeping it here for decoupling
+class GraniteMoeSharedMLP(nn.Module):
+    def __init__(
+        self,
+        config: GraniteMoeSharedConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.input_size = config.hidden_size
+        self.hidden_size = config.shared_intermediate_size
+        self.input_linear = MergedColumnParallelLinear(
+            input_size=self.input_size,
+            output_sizes=[self.hidden_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.input_linear",
+        )
+        self.output_linear = RowParallelLinear(
+            self.hidden_size,
+            self.input_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.output_linear",
+        )
+        if config.hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {config.hidden_act}. "
+                "Only silu is supported for now."
+            )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.input_linear(hidden_states)
+        x = self.act_fn(gate_up)
+        x, _ = self.output_linear(x)
+        return x
+
+
+class GraniteMoeHybridMambaDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: GraniteMoeHybridConfig,
+        layer_idx: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.residual_multiplier = config.residual_multiplier
+
+        self.mamba = MambaMixer2(
+            cache_params=config.mamba2_cache_params,
+            hidden_size=config.hidden_size,
+            use_conv_bias=config.mamba_conv_bias,
+            use_bias=config.mamba_proj_bias,
+            n_groups=config.mamba_n_groups,
+            rms_norm_eps=config.rms_norm_eps,
+            activation=config.hidden_act,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mixer",
+        )
+
+        self.block_sparse_moe = None
+        if getattr(config, "num_local_experts", 0) > 0:
+            self.block_sparse_moe = GraniteMoeMoE(
+                num_experts=config.num_local_experts,
+                top_k=config.num_experts_per_tok,
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                layer_id=layer_idx,
+                quant_config=quant_config,
+                tp_size=get_tensor_model_parallel_world_size(),
+                prefix=f"{prefix}.block_sparse_moe",
+            )
+
+        self.shared_mlp = (
+            None
+            if getattr(config, "shared_intermediate_size", 0) == 0
+            else GraniteMoeSharedMLP(
+                config, quant_config=quant_config, prefix=f"{prefix}.shared_mlp"
+            )
+        )
+
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+        forward_batch: ForwardBatch,
+    ):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        output = torch.empty_like(hidden_states)
+        attn_backend = forward_batch.attn_backend
+        assert isinstance(attn_backend, HybridLinearAttnBackend)
+        assert isinstance(attn_backend.linear_attn_backend, Mamba2AttnBackend)
+        attn_backend.linear_attn_backend.forward(
+            mixer=self.mamba,
+            layer_id=self.layer_idx,
+            hidden_states=hidden_states,
+            output=output,
+            use_triton_causal_conv=True,
+        )
+
+        hidden_states = residual + output * self.residual_multiplier
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        if self.shared_mlp is None:
+            if self.block_sparse_moe is not None:
+                hidden_states = self.block_sparse_moe(hidden_states)
+            # else: skip
+        else:
+            # create a copy since block_sparse_moe modifies in-place
+            if self.block_sparse_moe is not None:
+                moe_hidden_states = hidden_states.clone()
+                moe_hidden_states = self.block_sparse_moe(moe_hidden_states)
+                hidden_states = moe_hidden_states + self.shared_mlp(hidden_states)
+                del moe_hidden_states
+            else:
+                hidden_states = self.shared_mlp(hidden_states)
+        hidden_states = residual + hidden_states * self.residual_multiplier
+
+        return hidden_states, residual
+
+
+class GraniteMoeHybridAttention(nn.Module):
+    def __init__(
+        self,
+        config: GraniteMoeHybridConfig,
+        layer_id: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.causal = True
+        self.hidden_size = config.hidden_size
+        self.attention_bias = config.attention_bias
+        self.attention_multiplier = config.attention_multiplier
+        self.total_num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.total_num_heads
+        self.total_num_kv_heads = config.num_key_value_heads
+
+        # TensorParallel logic
+        tp_size = get_tensor_model_parallel_world_size()
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_key_value_heads = max(1, self.total_num_kv_heads // tp_size)
+
+        self.qkv_proj = QKVParallelLinear(
+            self.hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=self.attention_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            self.hidden_size,
+            self.hidden_size,
+            bias=self.attention_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        if config.position_embedding_type == "rope":
+
+            self.rotary_emb = get_rope(
+                head_size=self.head_dim,
+                rotary_dim=self.head_dim,  # its not in the config
+                max_position=config.max_position_embeddings,
+                base=config.rope_theta,
+                rope_scaling=config.rope_scaling,
+            )
+        else:
+            self.rotary_emb = None
+
+        self.attn = RadixAttention(
+            num_heads=self.num_heads,
+            head_dim=self.head_dim,
+            scaling=self.attention_multiplier,
+            num_kv_heads=self.num_key_value_heads,
+            layer_id=layer_id,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch | None = None,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        query, key, value = qkv.split(
+            [
+                self.num_heads * self.head_dim,
+                self.num_key_value_heads * self.head_dim,
+                self.num_key_value_heads * self.head_dim,
+            ],
+            dim=-1,
+        )
+
+        if self.rotary_emb is not None:
+            query, key = self.rotary_emb(positions, query, key)
+
+        hidden_states = self.attn(query, key, value, forward_batch=forward_batch)
+        del query, key, value
+
+        hidden_states = self.o_proj(hidden_states)[0]
+        return hidden_states
+
+
+class GraniteMoeHybridAttentionDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: GraniteMoeHybridConfig,
+        layer_idx: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.residual_multiplier = config.residual_multiplier
+
+        self.self_attn = GraniteMoeHybridAttention(
+            config,
+            layer_id=layer_idx,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+
+        self.block_sparse_moe = None
+        if getattr(config, "num_local_experts", 0) > 0:
+            self.block_sparse_moe = GraniteMoeMoE(
+                num_experts=config.num_local_experts,
+                top_k=config.num_experts_per_tok,
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                layer_id=layer_idx,
+                quant_config=quant_config,
+                tp_size=get_tensor_model_parallel_world_size(),
+                prefix=f"{prefix}.block_sparse_moe",
+            )
+
+        self.shared_mlp = (
+            None
+            if getattr(config, "shared_intermediate_size", 0) == 0
+            else GraniteMoeSharedMLP(
+                config, quant_config=quant_config, prefix=f"{prefix}.shared_mlp"
+            )
+        )
+
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+        forward_batch: ForwardBatch | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            forward_batch=forward_batch,
+        )
+        hidden_states = residual + hidden_states * self.residual_multiplier
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        if self.shared_mlp is None:
+            if self.block_sparse_moe is not None:
+                hidden_states = self.block_sparse_moe(hidden_states)
+            # else: skip
+        else:
+            # create a copy since block_sparse_moe modifies in-place
+            if self.block_sparse_moe is not None:
+                moe_hidden_states = hidden_states.clone()
+                moe_hidden_states = self.block_sparse_moe(moe_hidden_states)
+                hidden_states = moe_hidden_states + self.shared_mlp(hidden_states)
+                del moe_hidden_states
+            else:
+                hidden_states = self.shared_mlp(hidden_states)
+        hidden_states = residual + hidden_states * self.residual_multiplier
+
+        return hidden_states, residual
+
+
+ALL_DECODER_LAYER_TYPES = {
+    "attention": GraniteMoeHybridAttentionDecoderLayer,
+    "mamba": GraniteMoeHybridMambaDecoderLayer,
+}
+
+
+class GraniteMoeHybridModel(nn.Module):
+    def __init__(
+        self,
+        config: GraniteMoeHybridConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.config = config
+        self.quant_config = quant_config
+
+        self.vocab_size = config.vocab_size
+
+        self.pp_group = get_pp_group()
+
+        if self.pp_group.is_first_rank:
+            self.embed_tokens = VocabParallelEmbedding(
+                self.vocab_size,
+                config.hidden_size,
+                org_num_embeddings=config.vocab_size,
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        self.embedding_multiplier = config.embedding_multiplier
+
+        def get_layer(idx: int, prefix: str):
+            layer_idx = int(prefix.rsplit(".", 1)[1])
+            layer_class = ALL_DECODER_LAYER_TYPES[config.layer_types[layer_idx]]
+            return layer_class(
+                config,
+                layer_idx,
+                quant_config=quant_config,
+                prefix=prefix,
+            )
+
+        self.layers, self.start_layer, self.end_layer = make_layers(
+            config.num_hidden_layers,
+            get_layer,
+            pp_rank=self.pp_group.rank_in_group,
+            pp_size=self.pp_group.world_size,
+            prefix=f"{prefix}.layers",
+        )
+
+        if self.pp_group.is_last_rank:
+            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer(return_tuple=True)
+        self.layers_to_capture = []
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        """Get input embeddings from the model."""
+        return self.embed_tokens
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+    ) -> torch.Tensor:
+        if self.pp_group.is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_tokens(input_ids)
+                hidden_states = hidden_states * self.embedding_multiplier
+            residual = None
+        else:
+            assert pp_proxy_tensors is not None
+            hidden_states = pp_proxy_tensors["hidden_states"]
+            residual = pp_proxy_tensors["residual"]
+
+        aux_hidden_states = []
+        for i in range(self.start_layer, self.end_layer):
+            if i in self.layers_to_capture:
+                aux_hidden_states.append(hidden_states + residual)
+            layer = self.layers[i]
+            hidden_states, residual = layer(
+                positions,
+                hidden_states,
+                residual,
+                forward_batch,
+            )
+
+        if not self.pp_group.is_last_rank:
+            return PPProxyTensors(
+                {
+                    "hidden_states": hidden_states,
+                    "residual": residual,
+                }
+            )
+        else:
+            hidden_states, _ = self.norm(hidden_states, residual)
+
+        if len(aux_hidden_states) == 0:
+            return hidden_states
+
+        return hidden_states, aux_hidden_states
+
+
+class GraniteMoeHybridForCausalLM(
+    nn.Module,
+):
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+        "conv1d": ["conv1d"],
+        "in_proj": ["in_proj"],
+        "input_linear": ["input_linear"],
+    }
+    embedding_modules = {
+        "embed_tokens": "input_embeddings",
+        "lm_head": "output_embeddings",
+    }
+
+    def __init__(
+        self,
+        config: GraniteMoeHybridConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.capture_aux_hidden_states = False
+        self.pp_group = get_pp_group()
+
+        self.quant_config = quant_config
+        self.config = config
+        self.model = GraniteMoeHybridModel(
+            config=config,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "model"),
+        )
+
+        self.lm_head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=self.quant_config,
+            prefix=maybe_prefix(prefix, "lm_head"),
+        )
+
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+
+        self.logits_processor = LogitsProcessor(
+            config,
+            logit_scale=1 / self.config.logits_scaling,
+        )
+
+        self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
+
+    @property
+    def start_layer(self):
+        return self.model.start_layer
+
+    @property
+    def end_layer(self):
+        return self.model.end_layer
+
+    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: torch.Tensor = None,
+        get_embedding: bool = False,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+    ):
+        hidden_states = self.model(
+            input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
+        )
+
+        aux_hidden_states = None
+        if self.capture_aux_hidden_states:
+            hidden_states, aux_hidden_states = hidden_states
+
+        if self.pp_group.is_last_rank:
+            if not get_embedding:
+                return self.logits_processor(
+                    input_ids,
+                    hidden_states,
+                    self.lm_head,
+                    forward_batch,
+                    aux_hidden_states,
+                )
+            else:
+                return self.pooler(hidden_states, forward_batch)
+        else:
+            return hidden_states
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        # Params for weights, fp8 weight scales, fp8 activation scales
+        # (param_name, weight_name, expert_id, shard_id)
+        # layers.0.block_sparse_moe.expert_0.input_linear.input_scale
+        ckpt_gate_proj_name = "gate_proj"
+        ckpt_down_proj_name = "down_proj"
+        ckpt_up_proj_name = "up_proj"
+        num_experts = self.config.num_local_experts
+
+        return [
+            # (param_name, weight_name, expert_id, shard_id)
+            (
+                (
+                    "block_sparse_moe.experts.w13_"
+                    if weight_name in [ckpt_gate_proj_name, ckpt_up_proj_name]
+                    else "block_sparse_moe.experts.w2_"
+                ),
+                f"block_sparse_moe.experts.{expert_id}.{weight_name}.",
+                expert_id,
+                shard_id,
+            )
+            for expert_id in range(num_experts)
+            for shard_id, weight_name in [
+                ("w1", ckpt_gate_proj_name),
+                ("w2", ckpt_down_proj_name),
+                ("w3", ckpt_up_proj_name),
+            ]
+        ]
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        expert_params_mapping = self.get_expert_mapping()
+
+        def _load(n, p):
+            param = params_dict[n]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, p)
+            loaded_params.add(n)
+
+        def _load_shard(n, p, shard_id):
+            # Skip layers on other devices.
+            param = params_dict[n]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, p, shard_id)
+            loaded_params.add(n)
+
+        def _load_expert(n, p, name, shard_id, expert_id):
+            param = params_dict[n]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, p, name, shard_id=shard_id, expert_id=expert_id)
+            loaded_params.add(n)
+
+        def _load_quant_expert(name, loaded_weight):
+            for mapping in expert_params_mapping:
+                param_name, weight_name, expert_id, shard_id = mapping
+
+                if weight_name not in name:
+                    continue
+
+                name_mapped = name.replace(weight_name, param_name)
+
+                # Skip layers on other devices.
+                # if is_pp_missing_parameter(name_mapped, self):
+                #     continue
+
+                param = params_dict[name_mapped]
+                weight_loader = param.weight_loader
+                success = False
+
+                if weight_loader is not None:
+                    success = weight_loader(
+                        param,
+                        loaded_weight,
+                        name_mapped,
+                        shard_id=shard_id,
+                        expert_id=expert_id,
+                        return_success=True,
+                    )
+
+                if success:
+                    return name_mapped
+            return None
+
+        for n, p in weights:
+            if "A_log" in n:
+                n = n.replace("A_log", "A")
+
+            if self.quant_config is not None and (
+                scale_name := self.quant_config.get_cache_scale(n)
+            ):
+                # Loading kv cache quantization scales
+                loaded_weight = p
+                loaded_weight = (
+                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
+                )
+                _load(scale_name, loaded_weight)
+                loaded_params.add(scale_name)
+                continue
+
+            if _load_quant_expert(n, p):
+                continue
+
+            # Logic analogous to: https://github.com/vllm-project/vllm/blob/f49e5aff11c986ed4d45202b1716c5d74786efa9/vllm/model_executor/models/granitemoeshared.py#L215
+            # Mapping different experts' layout:
+            #  from HF (input_linear, output_linear, router)
+            #  to vLLM (experts_w13({e}.w1, {e}.w2), experts_w3({e}.w3), gate)
+            # The renaming and parameter loading logic is the same for weight
+            # and weight_scale tensors so we can reuse them without issues.
+            if n.endswith(".block_sparse_moe.input_linear.weight") or n.endswith(
+                ".block_sparse_moe.input_linear.weight_scale"
+            ):
+                for e in range(p.size(0)):
+                    w1_name = n.replace(
+                        ".block_sparse_moe.input_linear.weight",
+                        f".block_sparse_moe.experts.{e}.w1.weight",
+                    )
+                    w3_name = n.replace(
+                        ".block_sparse_moe.input_linear.weight",
+                        f".block_sparse_moe.experts.{e}.w3.weight",
+                    )
+                    w1_param, w3_param = p[e].chunk(2, dim=0)
+                    _load_expert(
+                        n.replace(".input_linear.", ".experts.w13_"),
+                        w1_param,
+                        w1_name,
+                        shard_id="w1",
+                        expert_id=e,
+                    )
+                    _load_expert(
+                        n.replace(".input_linear.", ".experts.w13_"),
+                        w3_param,
+                        w3_name,
+                        shard_id="w3",
+                        expert_id=e,
+                    )
+            elif n.endswith(".block_sparse_moe.output_linear.weight") or n.endswith(
+                ".block_sparse_moe.output_linear.weight_scale"
+            ):
+                for e in range(p.size(0)):
+                    w2_name = n.replace(
+                        ".block_sparse_moe.output_linear.weight",
+                        f".block_sparse_moe.experts.{e}.w2.weight",
+                    )
+                    w2_param = p[e]
+                    _load_expert(
+                        n.replace(".output_linear.", ".experts.w2_"),
+                        w2_param,
+                        w2_name,
+                        shard_id="w2",
+                        expert_id=e,
+                    )
+            elif n.endswith(".block_sparse_moe.router.layer.weight"):
+                gate_name = n.replace(
+                    ".block_sparse_moe.router.layer.weight",
+                    ".block_sparse_moe.gate.weight",
+                )
+                _load(gate_name, p)
+            else:
+                loaded = False
+                for param_name, weight_name, shard_id in stacked_params_mapping:
+                    if weight_name in n:
+                        _load_shard(
+                            n.replace(weight_name, param_name), p, shard_id=shard_id
+                        )
+                        loaded = True
+                if not loaded:
+                    _load(n, p)
+
+        return loaded_params
+
+
+EntryClass = [GraniteMoeHybridForCausalLM]

python/sglang/srt/server_args.py

@@ -1635,6 +1635,19 @@ class ServerArgs:
                 sm100_default_attention_backend="triton",
             )
 
+        elif model_arch == "GraniteMoeHybridForCausalLM":
+            hf_config = self.get_model_config().hf_config
+            has_mamba = any(
+                layer_type == "mamba"
+                for layer_type in getattr(hf_config, "layer_types", [])
+            )
+            if has_mamba:
+                self._handle_mamba_radix_cache(
+                    model_arch=model_arch,
+                    support_mamba_cache_extra_buffer=False,
+                    sm100_default_attention_backend="triton",
+                )
+
         elif model_arch in ["Lfm2ForCausalLM"]:
             self._handle_mamba_radix_cache(
                 model_arch=model_arch,

python/sglang/srt/utils/hf_transformers_utils.py

@@ -53,6 +53,7 @@ from sglang.srt.configs import (
     DotsVLMConfig,
     ExaoneConfig,
     FalconH1Config,
+    GraniteMoeHybridConfig,
     JetNemotronConfig,
     JetVLMConfig,
     KimiK25Config,
@@ -92,6 +93,7 @@ _CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [
     KimiLinearConfig,
     Qwen3NextConfig,
     FalconH1Config,
+    GraniteMoeHybridConfig,
     DotsVLMConfig,
     DotsOCRConfig,
     NemotronH_Nano_VL_V2_Config,