sglang 0.5.0rc1__py3-none-any.whl → 0.5.1__py3-none-any.whl

sglang/srt/models/interns1.py

@@ -4,8 +4,9 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.srt.distributed import parallel_state
+from sglang.srt.layers.attention import vision_utils
 from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
+from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.managers.mm_utils import (
     MultiModalityDataPaddingPatternTokenPairs,
@@ -20,6 +21,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.internvl import InternVisionModel
 from sglang.srt.models.qwen2 import Qwen2ForCausalLM
+from sglang.srt.models.qwen3 import Qwen3ForCausalLM
 from sglang.srt.models.qwen3_moe import Qwen3MoeForCausalLM
 from sglang.utils import logger
 
@@ -34,7 +36,7 @@ class InternS1ForConditionalGeneration(nn.Module):
         super().__init__()
         self.config = config
         self.quant_config = quant_config
-        self._update_hf_config()
+        vision_utils.update_vit_attn_dummy_heads_config(self.config)
         image_size = (
             getattr(config, "force_image_size", None) or config.vision_config.image_size
         )
@@ -69,6 +71,10 @@ class InternS1ForConditionalGeneration(nn.Module):
             self.language_model = Qwen3MoeForCausalLM(
                 config=config.text_config, quant_config=quant_config
             )
+        elif config.text_config.architectures[0] == "Qwen3ForCausalLM":
+            self.language_model = Qwen3ForCausalLM(
+                config=config.text_config, quant_config=quant_config
+            )
         else:
             raise NotImplementedError(
                 f"{config.text_config.architectures[0]} is not implemented."
@@ -86,21 +92,6 @@ class InternS1ForConditionalGeneration(nn.Module):
             nn.Linear(llm_hidden_size, llm_hidden_size),
         )
 
-    def _update_hf_config(self):
-        """update hf config to support tp"""
-        world_size = parallel_state.get_tensor_model_parallel_world_size()
-        num_heads = self.config.vision_config.num_attention_heads
-        head_dim = self.config.vision_config.hidden_size // num_heads
-        num_dummy_heads = 0
-
-        if num_heads % world_size != 0:
-            num_dummy_heads = (
-                (num_heads + world_size) // world_size
-            ) * world_size - num_heads
-
-        setattr(self.config.vision_config, "head_dim", head_dim)
-        setattr(self.config.vision_config, "num_dummy_heads", num_dummy_heads)
-
     def pixel_shuffle(self, x, scale_factor=0.5):
         n, w, h, c = x.size()
         # N, W, H, C --> N, W, H * scale, C // scale
@@ -183,34 +174,6 @@ class InternS1ForConditionalGeneration(nn.Module):
 
         return helper.pad_input_tokens(input_ids, mm_inputs)
 
-    def _pad_vit_attn_dummy_heads(self, name: str, loaded_weight: torch.Tensor):
-        """pad attn qkv weights for dummy heads"""
-        num_dummy_heads = self.config.vision_config.num_dummy_heads
-        if num_dummy_heads == 0:
-            return loaded_weight
-        head_dim = self.config.vision_config.head_dim
-
-        if any([_ in name for _ in ["attn.q_proj", "attn.k_proj", "attn.v_proj"]]):
-            if name.endswith(".weight"):
-                dummy_shape = [num_dummy_heads, head_dim, loaded_weight.shape[-1]]
-            elif name.endswith(".bias"):
-                dummy_shape = [num_dummy_heads, head_dim]
-            else:
-                raise RuntimeError(f"Unsupported weight with name={name}")
-            padded_weight = loaded_weight.new_zeros(dummy_shape)
-            loaded_weight = torch.cat(
-                [loaded_weight.unflatten(0, (-1, head_dim)), padded_weight], dim=0
-            ).flatten(0, 1)
-        if "attn.proj.weight" in name:
-            padded_weight = loaded_weight.new_zeros(
-                loaded_weight.shape[0], head_dim * num_dummy_heads
-            )
-            loaded_weight = torch.cat([loaded_weight, padded_weight], dim=-1)
-        if "attn.q_norm.weight" in name or "attn.k_norm.weight" in name:
-            padded_weight = loaded_weight.new_zeros(head_dim * num_dummy_heads)
-            loaded_weight = torch.cat([loaded_weight, padded_weight], dim=0)
-        return loaded_weight
-
     def _mapping_interns1_name(self, name):
         names_map = {
             "lm_head.weight": "language_model.lm_head.weight",
@@ -254,7 +217,7 @@ class InternS1ForConditionalGeneration(nn.Module):
         ]
         expert_params_mapping = []
         if "Qwen3MoeForCausalLM" in self.config.text_config.architectures:
-            expert_params_mapping = get_moe_impl_class().make_expert_params_mapping(
+            expert_params_mapping = FusedMoE.make_expert_params_mapping(
                 ckpt_gate_proj_name="gate_proj",
                 ckpt_down_proj_name="down_proj",
                 ckpt_up_proj_name="up_proj",
@@ -269,7 +232,9 @@ class InternS1ForConditionalGeneration(nn.Module):
                 continue
             name = self._mapping_interns1_name(name)
             if "vision_model" in name:
-                loaded_weight = self._pad_vit_attn_dummy_heads(name, loaded_weight)
+                loaded_weight = vision_utils.pad_vit_attn_dummy_heads(
+                    self.config, name, loaded_weight
+                )
 
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:

sglang/srt/models/internvl.py

@@ -10,9 +10,9 @@ from transformers import PretrainedConfig, PreTrainedModel
 from transformers.activations import ACT2FN
 from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
 
-from sglang.srt.distributed import parallel_state
+from sglang.srt.layers.attention import vision_utils
 from sglang.srt.layers.attention.vision import SingletonCache, VisionAttention
-from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
+from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.managers.mm_utils import (
     MultiModalityDataPaddingPatternTokenPairs,
@@ -412,7 +412,7 @@ class InternVLChatModel(nn.Module):
         super().__init__()
         self.config = config
         self.quant_config = quant_config
-        self._update_vision_config()
+        vision_utils.update_vit_attn_dummy_heads_config(self.config)
         image_size = config.force_image_size or config.vision_config.image_size
         patch_size = config.vision_config.patch_size
         self.patch_size = patch_size
@@ -462,21 +462,6 @@ class InternVLChatModel(nn.Module):
             nn.Linear(llm_hidden_size, llm_hidden_size),
         )
 
-    def _update_vision_config(self):
-        """update vision config to support tp"""
-        world_size = parallel_state.get_tensor_model_parallel_world_size()
-        num_heads = self.config.vision_config.num_attention_heads
-        head_dim = self.config.vision_config.hidden_size // num_heads
-        num_dummy_heads = 0
-
-        if num_heads % world_size != 0:
-            num_dummy_heads = (
-                (num_heads + world_size) // world_size
-            ) * world_size - num_heads
-
-        setattr(self.config.vision_config, "head_dim", head_dim)
-        setattr(self.config.vision_config, "num_dummy_heads", num_dummy_heads)
-
     def pixel_shuffle(self, x, scale_factor=0.5):
         n, w, h, c = x.size()
         # N, W, H, C --> N, W, H * scale, C // scale
@@ -559,36 +544,6 @@ class InternVLChatModel(nn.Module):
 
         return helper.pad_input_tokens(input_ids, mm_inputs)
 
-    def _pad_vit_attn_dummy_heads(self, name: str, loaded_weight: torch.Tensor):
-        """pad attn qkv weights for dummy heads"""
-        num_dummy_heads = self.config.vision_config.num_dummy_heads
-        if num_dummy_heads == 0:
-            return loaded_weight
-        head_dim = self.config.vision_config.head_dim
-
-        if "attn.qkv_proj" in name:
-            wq, wk, wv = loaded_weight.chunk(3, dim=0)
-            if name.endswith(".weight"):
-                dummy_shape = [num_dummy_heads, head_dim, wq.shape[-1]]
-            elif name.endswith(".bias"):
-                dummy_shape = [num_dummy_heads, head_dim]
-            else:
-                raise RuntimeError(f"Unsupported weight with name={name}")
-            pad_func = lambda x: torch.cat(
-                [x.unflatten(0, (-1, head_dim)), x.new_zeros(dummy_shape)], dim=0
-            ).flatten(0, 1)
-            wq, wk, wv = pad_func(wq), pad_func(wk), pad_func(wv)
-            loaded_weight = torch.cat([wq, wk, wv], dim=0)
-        if "attn.proj.weight" in name:
-            padded_weight = loaded_weight.new_zeros(
-                loaded_weight.shape[0], head_dim * num_dummy_heads
-            )
-            loaded_weight = torch.cat([loaded_weight, padded_weight], dim=-1)
-        if "attn.q_norm.weight" in name or "attn.k_norm.weight" in name:
-            padded_weight = loaded_weight.new_zeros(head_dim * num_dummy_heads)
-            loaded_weight = torch.cat([loaded_weight, padded_weight], dim=0)
-        return loaded_weight
-
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         expert_params_mapping = []
         if "InternLM2ForCausalLM" in self.config.llm_config.architectures:
@@ -616,7 +571,7 @@ class InternVLChatModel(nn.Module):
                 ("gate_up_proj", "up_proj", 1),
             ]
 
-            expert_params_mapping = get_moe_impl_class().make_expert_params_mapping(
+            expert_params_mapping = FusedMoE.make_expert_params_mapping(
                 ckpt_gate_proj_name="gate_proj",
                 ckpt_down_proj_name="down_proj",
                 ckpt_up_proj_name="up_proj",
@@ -699,8 +654,8 @@ class InternVLChatModel(nn.Module):
                             param, "weight_loader", default_weight_loader
                         )
                         if "vision_model" in name:
-                            loaded_weight = self._pad_vit_attn_dummy_heads(
-                                name, loaded_weight
+                            loaded_weight = vision_utils.pad_vit_attn_dummy_heads(
+                                self.config, name, loaded_weight
                             )
                         weight_loader(param, loaded_weight)
 

sglang/srt/models/llama.py

@@ -91,10 +91,18 @@ class LlamaMLP(nn.Module):
             )
         self.act_fn = SiluAndMul()
 
-    def forward(self, x, forward_batch=None):
+    def forward(
+        self,
+        x,
+        forward_batch=None,
+        use_reduce_scatter: bool = False,
+    ):
         gate_up, _ = self.gate_up_proj(x)
         x = self.act_fn(gate_up)
-        x, _ = self.down_proj(x)
+        x, _ = self.down_proj(
+            x,
+            skip_all_reduce=use_reduce_scatter,
+        )
         return x
 
 

sglang/srt/models/llama4.py

@@ -31,7 +31,7 @@ from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
 from sglang.srt.layers.dp_attention import (
     get_attention_tp_rank,
     get_attention_tp_size,
-    get_local_attention_dp_size,
+    is_dp_attention_enabled,
 )
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
@@ -45,7 +45,6 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.layers.rotary_embedding import get_rope
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import (
     ForwardBatch,
     ForwardMode,
@@ -131,14 +130,19 @@ class Llama4MoE(nn.Module):
             reduce_results=False,  # We need to do scatter before reduce
         )
 
-    def forward(self, hidden_states, forward_batch: ForwardBatch):
+    def forward(
+        self,
+        hidden_states,
+        forward_batch: ForwardBatch,
+        use_reduce_scatter: bool = False,
+    ):
         shared_out, routed_out = self._forward_core(
             hidden_states, forward_batch.forward_mode
         )
 
         out_aD = routed_out + shared_out
 
-        if self.tp_size > 1:
+        if self.tp_size > 1 and not use_reduce_scatter:
             out_aD = tensor_model_parallel_all_reduce(out_aD)
 
         return out_aD
@@ -359,7 +363,6 @@ class Llama4DecoderLayer(nn.Module):
         rope_theta = config.rope_theta
         rope_scaling = config.rope_scaling
         max_position_embeddings = config.max_position_embeddings
-        self.local_dp_size = get_local_attention_dp_size()
         self.attn_tp_size = get_attention_tp_size()
         self.attn_tp_rank = get_attention_tp_rank()
 
@@ -412,6 +415,7 @@ class Llama4DecoderLayer(nn.Module):
             layer_scatter_modes=self.layer_scatter_modes,
             input_layernorm=self.input_layernorm,
             post_attention_layernorm=self.post_attention_layernorm,
+            allow_reduce_scatter=True,
         )
 
     def _is_moe_layer(self, layer_id: int) -> bool:
@@ -441,8 +445,15 @@ class Llama4DecoderLayer(nn.Module):
             hidden_states, residual, forward_batch
         )
 
+        # For DP with padding, reduce scatter can be used instead of all-reduce.
+        use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
+            forward_batch
+        )
+
         # Fully Connected
-        hidden_states = self.feed_forward(hidden_states, forward_batch)
+        hidden_states = self.feed_forward(
+            hidden_states, forward_batch, use_reduce_scatter
+        )
         hidden_states, residual = self.layer_communicator.postprocess_layer(
             hidden_states, residual, forward_batch
         )
@@ -466,7 +477,7 @@ class Llama4Model(nn.Module):
             config.hidden_size,
             quant_config=quant_config,
             prefix=add_prefix("embed_tokens", prefix),
-            enable_tp=not global_server_args_dict["enable_dp_attention"],
+            enable_tp=not is_dp_attention_enabled(),
         )
         self.layers = make_layers(
             config.num_hidden_layers,

sglang/srt/models/minicpm3.py

@@ -37,7 +37,6 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-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, is_cuda

sglang/srt/models/mixtral.py

@@ -47,7 +47,6 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.managers.schedule_batch import global_server_args_dict
 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.utils import add_prefix, make_layers
@@ -104,7 +103,6 @@ class MixtralMoE(nn.Module):
             intermediate_size=intermediate_size,
             params_dtype=params_dtype,
             quant_config=quant_config,
-            tp_size=tp_size,
             prefix=add_prefix("experts", prefix),
         )