[PP] Add pp support for Qwen3-VL (#12333)

Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>
Signed-off-by: Kun(llfl) <i@imux.top>
Signed-off-by: Kun(llfl) <llfl@linux.alibaba.com>
Co-authored-by: kun-llfl <i@imux.top>
Co-authored-by: Kun(llfl) <llfl@linux.alibaba.com>

python/sglang/srt/managers/schedule_policy.py

@@ -455,6 +455,10 @@ class PrefillAdder:
 
     def add_chunked_req(self, req: Req):
         _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens))
+        # The chunked_req must be added to the list; otherwise, it will cause a memory leak.
+        # Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens.
+        if _rem_tokens <= 0:
+            _rem_tokens = self.rem_chunk_tokens
         truncated = req.extend_input_len > _rem_tokens
         req.extend_input_len = min(req.extend_input_len, _rem_tokens)
         req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]

python/sglang/srt/models/qwen3_omni_moe.py

@@ -614,7 +614,10 @@ class Qwen3OmniMoeForConditionalGeneration(PreTrainedModel):
                             and name_mapped not in params_dict
                         ):
                             continue
-                        param = params_dict[name_mapped]
+                        if name_mapped in params_dict.keys():
+                            param = params_dict[name_mapped]
+                        else:
+                            continue
                         # We should ask the weight loader to return success or
                         # not here since otherwise we may skip experts with
                         # # other available replicas.

python/sglang/srt/models/qwen3_vl.py

@@ -33,11 +33,13 @@ from sglang.srt.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
 )
+from sglang.srt.distributed.parallel_state import get_pp_group
 from sglang.srt.layers.attention.vision import VisionAttention
 from sglang.srt.layers.linear import ColumnParallelLinear, 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.utils import PPMissingLayer, get_layer_id
 from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
 from sglang.srt.managers.mm_utils import (
     MultiModalityDataPaddingPatternMultimodalTokens,
@@ -600,6 +602,7 @@ class Qwen3VLForConditionalGeneration(nn.Module):
         language_model_cls=Qwen3LLMModel,
     ) -> None:
         super().__init__()
+        self.pp_group = get_pp_group()
 
         self.use_data_parallel = get_global_server_args().mm_enable_dp_encoder
 
@@ -627,19 +630,22 @@ class Qwen3VLForConditionalGeneration(nn.Module):
                 quant_config=quant_config,
                 prefix=add_prefix("model", prefix),
             )
-
-            if self.config.tie_word_embeddings:
-                self.lm_head = self.model.embed_tokens
+            if self.pp_group.is_last_rank:
+                if self.pp_group.world_size == 1 and self.config.tie_word_embeddings:
+                    self.lm_head = self.model.embed_tokens
+                else:
+                    self.lm_head = ParallelLMHead(
+                        self.config.vocab_size,
+                        self.config.hidden_size,
+                        quant_config=quant_config,
+                        prefix=add_prefix("lm_head", prefix),
+                    )
             else:
-                self.lm_head = ParallelLMHead(
-                    self.config.vocab_size,
-                    self.config.hidden_size,
-                    quant_config=quant_config,
-                    prefix=add_prefix("lm_head", prefix),
-                )
+                self.lm_head = PPMissingLayer()
         else:
             # encoder_only mode: no language model, so no lm_head needed
             self.lm_head = None
+
         self.is_mrope_enabled = "mrope_section" in self.config.rope_scaling
 
         self.logits_processor = LogitsProcessor(self.config)
@@ -802,6 +808,7 @@ class Qwen3VLForConditionalGeneration(nn.Module):
         positions: torch.Tensor,
         forward_batch: ForwardBatch,
         get_embedding: bool = False,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ):
         """Run forward pass for Qwen3-VL.
 
@@ -835,14 +842,21 @@ class Qwen3VLForConditionalGeneration(nn.Module):
             multimodal_model=self,
             positions=positions,
             use_deepstack=self.use_deepstack,
+            pp_proxy_tensors=pp_proxy_tensors,
         )
 
-        if not get_embedding:
-            return self.logits_processor(
-                input_ids, hidden_states, self.lm_head, forward_batch
-            )
+        if self.pp_group.is_last_rank:
+            if not get_embedding:
+                return self.logits_processor(
+                    input_ids,
+                    hidden_states,
+                    self.lm_head,
+                    forward_batch,
+                )
+            else:
+                return self.pooler(hidden_states, forward_batch)
         else:
-            return self.pooler(hidden_states, forward_batch)
+            return hidden_states
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [
@@ -859,6 +873,27 @@ class Qwen3VLForConditionalGeneration(nn.Module):
                 continue
             if "language_model" in name:
                 name = name.replace(r"model.language_model.", r"model.")
+            layer_id = get_layer_id(name)
+
+            if self.pp_group.is_last_rank and "model.embed_tokens.weight" in name:
+                if "lm_head.weight" in params_dict:
+                    lm_head_param = params_dict["lm_head.weight"]
+                    weight_loader = getattr(
+                        lm_head_param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(lm_head_param, loaded_weight)
+
+            is_visual = "visual" in name
+            if (
+                not is_visual
+                and layer_id is not None
+                and hasattr(self.model, "start_layer")
+                and (
+                    layer_id < self.model.start_layer
+                    or layer_id >= self.model.end_layer
+                )
+            ):
+                continue
 
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
@@ -889,12 +924,11 @@ class Qwen3VLForConditionalGeneration(nn.Module):
                     # Skip loading extra bias for GPTQ models.
                     if name.endswith(".bias") and name not in params_dict:
                         continue
-                    # Skip loading visual/language model weights
-                    if (
-                        self.config.encoder_only or self.config.language_only
-                    ) and name not in params_dict:
+                    if name in params_dict.keys():
+                        param = params_dict[name]
+                    else:
                         continue
-                    param = params_dict[name]
+
                 except KeyError:
                     print(params_dict.keys())
                     raise

python/sglang/test/test_utils.py

@@ -59,6 +59,7 @@ DEFAULT_MODEL_NAME_FOR_TEST_MLA = "lmsys/sglang-ci-dsv3-test"
 DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN = "lmsys/sglang-ci-dsv3-test-NextN"
 
 # VL test models
+DEFAULT_MODEL_NAME_FOR_TEST_VL_PP = "Qwen/Qwen3-VL-2B-Thinking"
 DEFAULT_MODEL_NAME_FOR_TEST_GLM_41V_PP = "zai-org/GLM-4.1V-9B-Thinking"
 
 # NVFP4 models