diff --git a/docs/supported_models/text_generation/generative_models.md b/docs/supported_models/text_generation/generative_models.md index 1c74b1a5f..df5ff798b 100644 --- a/docs/supported_models/text_generation/generative_models.md +++ b/docs/supported_models/text_generation/generative_models.md @@ -68,3 +68,4 @@ in the GitHub search bar. | **Trinity** (Nano, Mini) | `arcee-ai/Trinity-Mini` | Arcee's foundational MoE Trinity family of models, open weights under Apache 2.0. | | **Falcon-H1** (0.5B–34B) | `tiiuae/Falcon-H1-34B-Instruct` | TII's hybrid Mamba-Transformer architecture combining attention and state-space models for efficient long-context inference. | | **Hunyuan-Large** (389B, MoE) | `tencent/Tencent-Hunyuan-Large` | Tencent's open-source MoE model with 389B total / 52B active parameters, featuring Cross-Layer Attention (CLA) for improved efficiency. | +| **IBM Granite 4.0 (Hybrid, Dense)** | `ibm-granite/granite-4.0-h-micro`, `ibm-granite/granite-4.0-micro` | IBM Granite 4.0 micro models: hybrid Mamba–MoE (`h-micro`) and dense (`micro`) variants. Enterprise-focused reasoning models | diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 965b4e305..3a3b37f54 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/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", diff --git a/python/sglang/srt/configs/granitemoehybrid.py b/python/sglang/srt/configs/granitemoehybrid.py new file mode 100644 index 000000000..8d2a7a7d1 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ab96f3f67..eb4a4c18a 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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 diff --git a/python/sglang/srt/models/granitemoehybrid.py b/python/sglang/srt/models/granitemoehybrid.py new file mode 100644 index 000000000..e18aeb466 --- /dev/null +++ b/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] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 79e0c51c2..1fb5463c9 100644 --- a/python/sglang/srt/server_args.py +++ b/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, diff --git a/python/sglang/srt/utils/hf_transformers_utils.py b/python/sglang/srt/utils/hf_transformers_utils.py index 41f70942d..9c6886fe5 100644 --- a/python/sglang/srt/utils/hf_transformers_utils.py +++ b/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, diff --git a/test/registered/models/test_generation_models.py b/test/registered/models/test_generation_models.py index d1041ea0d..2b4945d46 100644 --- a/test/registered/models/test_generation_models.py +++ b/test/registered/models/test_generation_models.py @@ -115,6 +115,10 @@ ALL_MODELS = [ "LiquidAI/LFM2.5-1.2B-Instruct", trust_remote_code=True, ), + ModelCase( + "ibm-granite/granite-4.0-h-micro", + trust_remote_code=True, + ), ] MAMBA_MODEL_PATHS = [