sglang 0.4.3.post2__py3-none-any.whl → 0.4.3.post4__py3-none-any.whl

sglang/srt/models/minicpmv.py

@@ -56,6 +56,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.utils import set_default_torch_dtype
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.qwen2 import Qwen2Config, Qwen2ForCausalLM
+from sglang.srt.utils import add_prefix
 
 RawImageType = Union[Image.Image, torch.Tensor]
 
@@ -158,14 +159,14 @@ class Idefics2VisionMLP(nn.Module):
             config.intermediate_size,
             bias=True,
             quant_config=quant_config,
-            prefix=f"{prefix}.fc1",
+            prefix=add_prefix("fc1", prefix),
         )
         self.fc2 = RowParallelLinear(
             config.intermediate_size,
             config.hidden_size,
             bias=True,
             quant_config=quant_config,
-            prefix=f"{prefix}.fc2",
+            prefix=add_prefix("fc2", prefix),
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -199,10 +200,14 @@ class Idefics2EncoderLayer(nn.Module):
             use_context_forward=False,
             use_full_precision_softmax=True,
             flatten_batch=False,
-            prefix=f"{prefix}.self_attn",
+            prefix=add_prefix("self_attn", prefix),
         )
         self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
-        self.mlp = Idefics2VisionMLP(config, quant_config=quant_config)
+        self.mlp = Idefics2VisionMLP(
+            config,
+            quant_config=quant_config,
+            prefix=add_prefix("mlp", prefix),
+        )
         self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
 
     def forward(
@@ -242,6 +247,7 @@ class Idefics2Encoder(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
 
@@ -251,8 +257,9 @@ class Idefics2Encoder(nn.Module):
                 Idefics2EncoderLayer(
                     config,
                     quant_config=quant_config,
+                    prefix=add_prefix(f"layers.{i}", prefix),
                 )
-                for _ in range(config.num_hidden_layers)
+                for i in range(config.num_hidden_layers)
             ]
         )
 
@@ -379,13 +386,18 @@ class Idefics2VisionTransformer(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
 
         embed_dim = config.hidden_size
         self.config = config
         self.embeddings = Idefics2VisionEmbeddings(config)
-        self.encoder = Idefics2Encoder(config=config, quant_config=quant_config)
+        self.encoder = Idefics2Encoder(
+            config=config,
+            quant_config=quant_config,
+            prefix=add_prefix("encoder", prefix),
+        )
         self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
 
     def get_input_embeddings(self):
@@ -503,7 +515,7 @@ class BaseResampler(nn.Module):
                 embed_dim,
                 bias=False,
                 quant_config=quant_config,
-                prefix=f"{prefix}.kv_proj",
+                prefix=add_prefix("kv_proj", prefix),
             )
         else:
             # Maintain the same return value with ReplicatedLinear.forward
@@ -660,6 +672,7 @@ class MiniCPMVBaseModel(nn.Module):
         *,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         # All MiniCPM-V models disable `tie_word_embeddings` but
@@ -669,8 +682,12 @@ class MiniCPMVBaseModel(nn.Module):
         self.config = config
 
         self.version = get_version_by_config(self.config)
-        self.llm = self.init_llm(config=config, quant_config=quant_config)
-        self.vpm = self.init_vision_module(config, quant_config)
+        self.llm = self.init_llm(
+            config=config, quant_config=quant_config, prefix=add_prefix("llm", prefix)
+        )
+        self.vpm = self.init_vision_module(
+            config, quant_config, add_prefix("vpm", prefix)
+        )
         self.vision_dim = (
             self.vpm.embed_dim
             if self.version == (2, 0)
@@ -679,7 +696,10 @@ class MiniCPMVBaseModel(nn.Module):
         self.embed_dim = self.config.hidden_size
 
         self.resampler = self.init_resampler(
-            self.embed_dim, self.vision_dim, quant_config=quant_config
+            self.embed_dim,
+            self.vision_dim,
+            quant_config=quant_config,
+            prefix=add_prefix("resampler", prefix),
         )
 
         self.logits_processor = LogitsProcessor(config)
@@ -937,6 +957,7 @@ class MiniCPMVBaseModel(nn.Module):
         self,
         config: Qwen2Config,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> nn.Module:
         raise NotImplementedError
 
@@ -944,6 +965,7 @@ class MiniCPMVBaseModel(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig],
+        prefix: str = "",
     ) -> nn.Module:
         raise NotImplementedError
 
@@ -952,6 +974,7 @@ class MiniCPMVBaseModel(nn.Module):
         embed_dim: int,
         vision_dim: int,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> nn.Module:
         raise NotImplementedError
 
@@ -1011,24 +1034,27 @@ class MiniCPMV2_6(MiniCPMVBaseModel):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
-        super().__init__(config=config, quant_config=quant_config)
+        super().__init__(config=config, quant_config=quant_config, prefix=prefix)
         assert self.version == (2, 6)
 
     def init_llm(
         self,
         config: Qwen2Config,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> nn.Module:
-        return Qwen2ForCausalLM(config=config, quant_config=quant_config)
+        return Qwen2ForCausalLM(config=config, quant_config=quant_config, prefix=prefix)
 
     def init_vision_module(
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig],
+        prefix: str = "",
     ) -> nn.Module:
         model = Idefics2VisionTransformer(
-            config=config.vision_config, quant_config=quant_config
+            config=config.vision_config, quant_config=quant_config, prefix=prefix
         )
         if self.config.drop_vision_last_layer:
             model.encoder.layers = model.encoder.layers[:-1]
@@ -1042,6 +1068,7 @@ class MiniCPMV2_6(MiniCPMVBaseModel):
         embed_dim: int,
         vision_dim: int,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> nn.Module:
         with set_default_torch_dtype(torch.float16):
             # The resampler in 2.6 remains consistent with the one in 2.5.
@@ -1051,6 +1078,7 @@ class MiniCPMV2_6(MiniCPMVBaseModel):
                 num_heads=embed_dim // 128,
                 kv_dim=vision_dim,
                 quant_config=quant_config,
+                prefix=prefix,
             )
 
         return resampler.to(device="cuda", dtype=torch.get_default_dtype())
@@ -1207,6 +1235,7 @@ class MiniCPMV:
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
 
@@ -1221,7 +1250,9 @@ class MiniCPMV:
             raise ValueError("Currently, MiniCPMV only supports versions 2.6")
 
         try:
-            minicpmv = instance_class(config=config, quant_config=quant_config)
+            minicpmv = instance_class(
+                config=config, quant_config=quant_config, prefix=prefix
+            )
             self.minicpmv = minicpmv
         except Exception as e:
             print(f"Failed to instantiate MiniCPMV: {e}")

sglang/srt/models/mixtral.py

@@ -45,6 +45,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 from sglang.srt.managers.schedule_batch import global_server_args_dict
 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
 
 
 class MixtralMoE(nn.Module):
@@ -78,7 +79,7 @@ class MixtralMoE(nn.Module):
             bias=False,
             params_dtype=params_dtype,
             quant_config=None,
-            prefix=f"{prefix}.gate",
+            prefix=add_prefix("gate", prefix),
         )
         MoEImpl = EPMoE if global_server_args_dict["enable_ep_moe"] else FusedMoE
         self.experts = MoEImpl(
@@ -90,7 +91,7 @@ class MixtralMoE(nn.Module):
             renormalize=True,
             quant_config=quant_config,
             tp_size=tp_size,
-            prefix=f"{prefix}.experts",
+            prefix=add_prefix("experts", prefix),
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -146,14 +147,14 @@ class MixtralAttention(nn.Module):
             self.total_num_kv_heads,
             bias=False,
             quant_config=quant_config,
-            prefix=f"{prefix}.qkv_proj",
+            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=f"{prefix}.o_proj",
+            prefix=add_prefix("o_proj", prefix),
         )
         self.rotary_emb = get_rope(
             self.head_dim,
@@ -168,6 +169,7 @@ class MixtralAttention(nn.Module):
             self.scaling,
             num_kv_heads=self.num_kv_heads,
             layer_id=layer_id,
+            prefix=add_prefix("attn", prefix),
         )
 
     def forward(
@@ -204,7 +206,7 @@ class MixtralDecoderLayer(nn.Module):
             layer_id=layer_id,
             rope_theta=rope_theta,
             quant_config=quant_config,
-            prefix=f"{prefix}.self_attn",
+            prefix=add_prefix("self_attn", prefix),
         )
         self.block_sparse_moe = MixtralMoE(
             num_experts=config.num_local_experts,
@@ -212,7 +214,7 @@ class MixtralDecoderLayer(nn.Module):
             hidden_size=config.hidden_size,
             intermediate_size=config.intermediate_size,
             quant_config=quant_config,
-            prefix=f"{prefix}.block_sparse_moe",
+            prefix=add_prefix("block_sparse_moe", prefix),
         )
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(
@@ -258,11 +260,15 @@ class MixtralModel(nn.Module):
         self.embed_tokens = VocabParallelEmbedding(
             config.vocab_size,
             config.hidden_size,
+            prefix=add_prefix("embed_tokens", prefix),
         )
         self.layers = nn.ModuleList(
             [
                 MixtralDecoderLayer(
-                    config, i, quant_config=quant_config, prefix=f"{prefix}.layers"
+                    config,
+                    i,
+                    quant_config=quant_config,
+                    prefix=add_prefix(f"layers.{i}", prefix),
                 )
                 for i in range(config.num_hidden_layers)
             ]
@@ -296,12 +302,17 @@ class MixtralForCausalLM(nn.Module):
         self,
         config: MixtralConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.config = config
         self.quant_config = quant_config
-        self.model = MixtralModel(config, quant_config=quant_config, prefix="model")
-        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
+        self.model = MixtralModel(
+            config, quant_config=quant_config, prefix=add_prefix("model", prefix)
+        )
+        self.lm_head = ParallelLMHead(
+            config.vocab_size, config.hidden_size, prefix=add_prefix("lm_head", prefix)
+        )
         self.logits_processor = LogitsProcessor(config)
 
     def forward(

sglang/srt/models/mixtral_quant.py

@@ -45,6 +45,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 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
 
 
 class MixtralMLP(nn.Module):
@@ -54,6 +55,7 @@ class MixtralMLP(nn.Module):
         hidden_size: int,
         intermediate_size: int,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.num_experts = num_experts
@@ -61,13 +63,25 @@ class MixtralMLP(nn.Module):
         self.hidden_dim = hidden_size
 
         self.w1 = ReplicatedLinear(
-            self.hidden_dim, self.ffn_dim, bias=False, quant_config=quant_config
+            self.hidden_dim,
+            self.ffn_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("w1", prefix),
         )
         self.w2 = ReplicatedLinear(
-            self.ffn_dim, self.hidden_dim, bias=False, quant_config=quant_config
+            self.ffn_dim,
+            self.hidden_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("w2", prefix),
         )
         self.w3 = ReplicatedLinear(
-            self.hidden_dim, self.ffn_dim, bias=False, quant_config=quant_config
+            self.hidden_dim,
+            self.ffn_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("w3", prefix),
         )
 
         # TODO: Use vllm's SiluAndMul
@@ -87,6 +101,7 @@ class MixtralMoE(nn.Module):
         self,
         config: MixtralConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.config = config
@@ -114,6 +129,7 @@ class MixtralMoE(nn.Module):
                         config.hidden_size,
                         config.intermediate_size,
                         quant_config=quant_config,
+                        prefix=add_prefix(f"experts.{idx}", prefix),
                     )
                     if idx in self.expert_indicies
                     else None
@@ -122,7 +138,11 @@ class MixtralMoE(nn.Module):
             ]
         )
         self.gate = ReplicatedLinear(
-            config.hidden_size, self.num_total_experts, bias=False, quant_config=None
+            config.hidden_size,
+            self.num_total_experts,
+            bias=False,
+            quant_config=None,
+            prefix=add_prefix("gate", prefix),
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -159,6 +179,7 @@ class MixtralAttention(nn.Module):
         max_position: int = 4096 * 32,
         rope_theta: float = 10000,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
@@ -189,12 +210,14 @@ class MixtralAttention(nn.Module):
             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.rotary_emb = get_rope(
             self.head_dim,
@@ -209,6 +232,7 @@ class MixtralAttention(nn.Module):
             self.scaling,
             num_kv_heads=self.num_kv_heads,
             layer_id=layer_id,
+            prefix=add_prefix("attn", prefix),
         )
 
     def forward(
@@ -231,6 +255,7 @@ class MixtralDecoderLayer(nn.Module):
         config: MixtralConfig,
         layer_id: int = 0,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -244,8 +269,13 @@ class MixtralDecoderLayer(nn.Module):
             layer_id=layer_id,
             rope_theta=rope_theta,
             quant_config=quant_config,
+            prefix=add_prefix("self_attn", prefix),
+        )
+        self.block_sparse_moe = MixtralMoE(
+            config=config,
+            quant_config=quant_config,
+            prefix=add_prefix("block_sparse_moe", prefix),
         )
-        self.block_sparse_moe = MixtralMoE(config=config, quant_config=quant_config)
         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
@@ -281,6 +311,7 @@ class MixtralModel(nn.Module):
         self,
         config: MixtralConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.padding_idx = config.pad_token_id
@@ -289,10 +320,16 @@ class MixtralModel(nn.Module):
         self.embed_tokens = VocabParallelEmbedding(
             config.vocab_size,
             config.hidden_size,
+            prefix=add_prefix("embed_tokens", prefix),
         )
         self.layers = nn.ModuleList(
             [
-                MixtralDecoderLayer(config, i, quant_config=quant_config)
+                MixtralDecoderLayer(
+                    config,
+                    i,
+                    quant_config=quant_config,
+                    prefix=add_prefix(f"layers.{i}", prefix),
+                )
                 for i in range(config.num_hidden_layers)
             ]
         )
@@ -324,12 +361,17 @@ class QuantMixtralForCausalLM(nn.Module):
         self,
         config: MixtralConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.config = config
         self.quant_config = quant_config
-        self.model = MixtralModel(config, quant_config=quant_config)
-        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
+        self.model = MixtralModel(
+            config, quant_config=quant_config, prefix=add_prefix("model", prefix)
+        )
+        self.lm_head = ParallelLMHead(
+            config.vocab_size, config.hidden_size, prefix=add_prefix("lm_head", prefix)
+        )
         self.logits_processor = LogitsProcessor(config)
 
     @torch.no_grad()