Support Gemma4 Pipeline Parallelism (#25284)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/srt/models/gemma4_causal.py

@@ -14,7 +14,7 @@
 
 import logging
 import re
-from typing import Iterable, List, Optional, Set, Tuple
+from typing import Iterable, List, Optional, Set, Tuple, Union
 
 import torch
 from torch import nn
@@ -25,6 +25,7 @@ from transformers import (
 )
 
 from sglang.srt.distributed import (
+    get_pp_group,
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
 )
@@ -45,8 +46,9 @@ from sglang.srt.layers.moe.topk import TopK
 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, get_layer_id
 from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
@@ -68,6 +70,59 @@ Gemma4MLP = Gemma3MLP
 Gemma4TextScaledWordEmbedding = Gemma3TextScaledWordEmbedding
 
 
+def pp_filter_load_weight(
+    name,
+    loaded_weight,
+    *,
+    pp_group,
+    start_layer,
+    end_layer,
+    params_dict,
+    loaded_params,
+    tie_word_embeddings,
+    embed_weight_name,
+    first_rank_only_patterns=(),
+    last_rank_only_prefixes=(),
+    head_param_name="lm_head.weight",
+):
+    """Shared PP filter for Gemma4 load_weights paths.
+
+    Returns True if the caller should ``continue`` (handled or skipped),
+    False otherwise.  No-op when ``pp_group.world_size == 1``.
+
+    Handles three concerns in order:
+      1. Drop transformer-layer weights outside [start_layer, end_layer).
+      2. Route the tied ``embed_tokens.weight`` to ``lm_head`` on the last
+         rank (under PP, embed and lm_head live on different ranks so they
+         can't be tied via module aliasing).
+      3. Skip rank-local module weights on the wrong rank.
+    """
+    if pp_group.world_size <= 1:
+        return False
+
+    layer_id = get_layer_id(name)
+    if layer_id is not None and (layer_id < start_layer or layer_id >= end_layer):
+        return True
+
+    if tie_word_embeddings and pp_group.is_last_rank and name == embed_weight_name:
+        head_param = params_dict.get(head_param_name)
+        if head_param is not None:
+            wl = getattr(head_param, "weight_loader", default_weight_loader)
+            wl(head_param, loaded_weight)
+            loaded_params.add(head_param_name)
+        return True
+
+    if not pp_group.is_first_rank and any(p in name for p in first_rank_only_patterns):
+        return True
+
+    if not pp_group.is_last_rank and any(
+        name.startswith(p) for p in last_rank_only_prefixes
+    ):
+        return True
+
+    return False
+
+
 class Gemma4Router(nn.Module):
     """Router for Gemma4 MoE that preprocesses input before projection.
 
@@ -673,15 +728,12 @@ class Gemma4TextModel(PreTrainedModel):
         self.quant_config = quant_config
         self.vocab_size = config.vocab_size
         self.padding_idx = getattr(config, "pad_token_id", None)
+        self.pp_group = get_pp_group()
 
-        self.embed_tokens = Gemma4TextScaledWordEmbedding(
+        # Token / per-layer embedding tables and the per-layer projection only
-            config.vocab_size,
+        # produce activations consumed at the model entry, so they live on the
-            config.hidden_size,
+        # first PP rank only.  Other ranks substitute PPMissingLayer so that
-            self.padding_idx,
+        # parameter iteration still works (load_weights skips them explicitly).
-            embed_scale=self.config.hidden_size**0.5,  # embedded normalizer
-        )
-
-        # Per-layer input embeddings
         self.hidden_size = config.hidden_size
         self.hidden_size_per_layer_input = (
             getattr(config, "hidden_size_per_layer_input", None) or 0
@@ -690,7 +742,43 @@ class Gemma4TextModel(PreTrainedModel):
             getattr(config, "vocab_size_per_layer_input", None) or config.vocab_size
         )
 
-        if self.hidden_size_per_layer_input and self.hidden_size_per_layer_input > 0:
+        # PLE-enabled variants (E2B/E4B) forward `per_layer_inputs` through
+        # the PP proxy, but cuda_graph_runner hardcodes the proxy schema to
+        # {hidden_states, residual} and silently drops any extra keys at
+        # replay time.  Empirically this corrupts E4B output to garbage on
+        # non-first PP ranks (eager path produces correct output and
+        # GSM8K ~0.92, cuda-graph path emits token soup).  Refuse the
+        # combination until the runner becomes schema-aware; users can run
+        # PP + PLE eagerly with --disable-cuda-graph.
+        if self.pp_group.world_size > 1 and self.hidden_size_per_layer_input > 0:
+            sa = get_global_server_args()
+            if sa is not None and not sa.disable_cuda_graph:
+                raise ValueError(
+                    "Pipeline parallelism is currently incompatible with "
+                    "per-layer-input (PLE) embeddings under CUDA graph: "
+                    "the runner's PP proxy schema is hardcoded to "
+                    "{hidden_states, residual} and silently drops "
+                    "per_layer_inputs, corrupting per-layer contributions on "
+                    "non-first PP ranks. Workarounds: (a) pass "
+                    "--disable-cuda-graph to fall back to eager replay, or "
+                    "(b) use tensor parallelism (--tp-size) instead of PP."
+                )
+
+        if self.pp_group.is_first_rank:
+            self.embed_tokens = Gemma4TextScaledWordEmbedding(
+                config.vocab_size,
+                config.hidden_size,
+                self.padding_idx,
+                embed_scale=self.config.hidden_size**0.5,  # embedded normalizer
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        if (
+            self.pp_group.is_first_rank
+            and self.hidden_size_per_layer_input
+            and self.hidden_size_per_layer_input > 0
+        ):
             self.embed_tokens_per_layer = Gemma4TextScaledWordEmbedding(
                 self.vocab_size_per_layer_input,
                 config.num_hidden_layers * self.hidden_size_per_layer_input,
@@ -721,7 +809,7 @@ class Gemma4TextModel(PreTrainedModel):
             self.per_layer_input_scale = None
             self.per_layer_projection_scale = None
 
-        self.layers = make_layers(
+        self.layers, self.start_layer, self.end_layer = make_layers(
             config.num_hidden_layers,
             lambda idx, prefix: Gemma4DecoderLayer(
                 layer_id=idx,
@@ -729,10 +817,15 @@ class Gemma4TextModel(PreTrainedModel):
                 quant_config=quant_config,
                 prefix=prefix,
             ),
+            pp_rank=self.pp_group.rank_in_group,
+            pp_size=self.pp_group.world_size,
             prefix=add_prefix("layers", prefix),
         )
 
-        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        if self.pp_group.is_last_rank:
+            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
         self.layers_to_capture = []
         self.post_init()
 
@@ -817,24 +910,35 @@ class Gemma4TextModel(PreTrainedModel):
         forward_batch: ForwardBatch,
         input_embeds: torch.Tensor = None,
         per_layer_inputs: Optional[torch.Tensor] = None,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]], PPProxyTensors]:
-        if (input_ids is None) ^ (input_embeds is not None):
+        if self.pp_group.is_first_rank:
-            raise ValueError(
+            if (input_ids is None) ^ (input_embeds is not None):
-                "You must specify exactly one of input_ids or inputs_embeds"
+                raise ValueError(
+                    "You must specify exactly one of input_ids or inputs_embeds"
+                )
+
+            if input_ids is not None:
+                input_embeds = self.embed_tokens(input_ids)
+                per_layer_inputs = self.get_per_layer_inputs(input_ids)
+            per_layer_inputs = self.project_per_layer_inputs(
+                input_embeds, per_layer_inputs
             )
-
+            hidden_states = input_embeds
-        if input_ids is not None:
+        else:
-            input_embeds = self.embed_tokens(input_ids)
+            assert (
-            per_layer_inputs = self.get_per_layer_inputs(input_ids)
+                pp_proxy_tensors is not None
-        per_layer_inputs = self.project_per_layer_inputs(input_embeds, per_layer_inputs)
+            ), "pp_proxy_tensors is required on non-first PP ranks"
-
+            hidden_states = pp_proxy_tensors["hidden_states"]
-        hidden_states = input_embeds
+            # PLE inputs were computed on rank 0 and forwarded along the
+            # pipeline; non-PLE models simply omit the key.
+            per_layer_inputs = pp_proxy_tensors.tensors.get("per_layer_inputs", None)
 
         aux_hidden_states = []
-        num_layers = len(self.layers)
+        num_layers = self.config.num_hidden_layers
 
-        for layer_idx, layer in enumerate(self.layers):
+        for layer_idx in range(self.start_layer, self.end_layer):
             if layer_idx in self.layers_to_capture:
                 aux_hidden_states.append(hidden_states)
 
@@ -842,6 +946,7 @@ class Gemma4TextModel(PreTrainedModel):
                 per_layer_input = per_layer_inputs[:, layer_idx, :]
             else:
                 per_layer_input = None
+            layer = self.layers[layer_idx]
             layer_outputs = layer(
                 positions=positions,
                 hidden_states=hidden_states,
@@ -850,7 +955,23 @@ class Gemma4TextModel(PreTrainedModel):
                 **kwargs,
             )
             hidden_states = layer_outputs[0]
-            residual = layer_outputs[1] if len(layer_outputs) > 1 else None
+            # Gemma4DecoderLayer.forward always returns (hidden_states, None);
+            # the residual is fused inside the layer, so nothing to thread.
+
+        if not self.pp_group.is_last_rank:
+            # cuda_graph_runner allocates a fixed PP-proxy schema of
+            # {hidden_states, residual} and KeyErrors if a model omits a key.
+            # Gemma4 fuses the residual inside each layer so we don't have a
+            # standalone tensor to forward; emit a zero placeholder instead so
+            # graph replay can still copy it.  The receiving stage never reads
+            # this key.
+            proxy = {
+                "hidden_states": hidden_states,
+                "residual": torch.zeros_like(hidden_states),
+            }
+            if per_layer_inputs is not None:
+                proxy["per_layer_inputs"] = per_layer_inputs
+            return PPProxyTensors(proxy)
 
         # Capture the output of the last layer if requested.
         # layers_to_capture uses +1 offset, so num_layers means
@@ -858,10 +979,7 @@ class Gemma4TextModel(PreTrainedModel):
         if num_layers in self.layers_to_capture:
             aux_hidden_states.append(hidden_states)
 
-        if residual is None:
+        hidden_states = self.norm(hidden_states)
-            hidden_states = self.norm(hidden_states)
-        else:
-            hidden_states, _ = self.norm(hidden_states, residual)
 
         if len(aux_hidden_states) == 0:
             return hidden_states
@@ -918,25 +1036,46 @@ class Gemma4ForCausalLM(PreTrainedModel):
         prefix: str = "",
     ) -> None:
         super().__init__(config=config)
+        self.pp_group = get_pp_group()
         self.config = config
         self.quant_config = quant_config
+
         self.model = Gemma4TextModel(
             config=config, quant_config=quant_config, prefix=add_prefix("model", prefix)
         )
         self.logits_processor = LogitsProcessor(config)
 
-        if self.config.tie_word_embeddings:
+        # tie_word_embeddings ties lm_head to embed_tokens, but with PP those
+        # tensors live on opposite ranks (first vs last).  In the PP > 1 case
+        # we materialize a real ParallelLMHead on the last rank and route the
+        # checkpoint's embed_tokens.weight into it during load_weights.
+        if self.pp_group.world_size == 1 and self.config.tie_word_embeddings:
             self.lm_head = self.model.embed_tokens
-        else:
+        elif self.pp_group.is_last_rank:
             self.lm_head = ParallelLMHead(
                 config.vocab_size,
                 config.hidden_size,
                 quant_config=quant_config,
                 prefix=add_prefix("lm_head", prefix),
             )
+        else:
+            self.lm_head = PPMissingLayer()
+
         self.capture_aux_hidden_states = False
         self.post_init()
 
+    def tie_weights(self, *args, **kwargs):
+        # HF's PreTrainedModel.tie_weights uses ``_tied_weights_keys`` to bind
+        # ``lm_head.weight`` to ``model.embed_tokens.weight``.  Under PP those
+        # tensors live on different ranks (embed on first, head on last) and
+        # the missing side is a PPMissingLayer with no ``weight`` attribute,
+        # which makes the default tie_weights crash.  load_weights routes the
+        # checkpoint embedding into lm_head explicitly, so the tie is a no-op
+        # here when PP is active.
+        if self.pp_group.world_size > 1:
+            return
+        super().tie_weights(*args, **kwargs)
+
     def get_input_embeddings(self) -> nn.Embedding:
         return self.model.embed_tokens
 
@@ -957,17 +1096,24 @@ class Gemma4ForCausalLM(PreTrainedModel):
         forward_batch: ForwardBatch,
         input_embeds: torch.Tensor = None,
         per_layer_inputs: Optional[torch.Tensor] = None,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
         **kwargs,
-    ) -> LogitsProcessor:
+    ) -> Union[LogitsProcessor, PPProxyTensors]:
         hidden_states = self.model(
             input_ids,
             positions,
             forward_batch,
             input_embeds,
             per_layer_inputs,
+            pp_proxy_tensors=pp_proxy_tensors,
             **kwargs,
         )
 
+        if not self.pp_group.is_last_rank:
+            # `hidden_states` here is actually a PPProxyTensors handed off to
+            # the next stage; logits processing only happens on the last rank.
+            return hidden_states
+
         aux_hidden_states = None
         if self.capture_aux_hidden_states:
             hidden_states, aux_hidden_states = hidden_states
@@ -1022,6 +1168,25 @@ class Gemma4ForCausalLM(PreTrainedModel):
             if ".experts." in name and ".moe.experts." not in name:
                 name = name.replace(".experts.", ".moe.experts.")
 
+            if pp_filter_load_weight(
+                name,
+                loaded_weight,
+                pp_group=self.pp_group,
+                start_layer=self.model.start_layer,
+                end_layer=self.model.end_layer,
+                params_dict=params_dict,
+                loaded_params=loaded_params,
+                tie_word_embeddings=self.config.tie_word_embeddings,
+                embed_weight_name="model.embed_tokens.weight",
+                first_rank_only_patterns=(
+                    "embed_tokens",
+                    "per_layer_model_projection",
+                    "per_layer_projection_norm",
+                ),
+                last_rank_only_prefixes=("model.norm.", "lm_head."),
+            ):
+                continue
+
             # attention_k_eq_v: full-attention layers have no v_proj in the
             # checkpoint (K and V share weights).  When we see a k_proj weight
             # for one of these layers, load it into both the "k" and "v" shards
@@ -1120,6 +1285,19 @@ class Gemma4ForCausalLM(PreTrainedModel):
         return self._shard_weight(self.model.embed_tokens.weight)
 
     def get_embed_and_head(self):
+        if self.pp_group.world_size > 1:
+            # Under PP, embed_tokens lives on the first rank and lm_head on
+            # the last; neither rank holds both tensors, so we can't return
+            # the pair locally without a cross-stage gather.  Callers (RL
+            # weight sync, remote weight loader) currently assume a
+            # single-rank view — fail loudly rather than dereference a
+            # PPMissingLayer.
+            raise NotImplementedError(
+                "get_embed_and_head() is not implemented for Gemma4ForCausalLM "
+                "under pipeline parallelism. embed_tokens lives on the first "
+                "PP rank and lm_head on the last; use --pp-size 1 if you "
+                "need this API."
+            )
         embed = self._shard_weight(self.model.embed_tokens.weight)
         head = self._shard_weight(self.lm_head.weight)
         return embed, head

python/sglang/srt/models/gemma4_mm.py

@@ -28,11 +28,14 @@ from transformers import (
     PreTrainedModel,
 )
 
+from sglang.srt.distributed import get_pp_group
 from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
 from sglang.srt.layers.layernorm import Gemma4RMSNorm
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.utils import PPMissingLayer
+from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
 from sglang.srt.managers.mm_utils import (
     MultiModalityDataPaddingPatternMultimodalTokens,
     general_mm_embed_routine,
@@ -43,13 +46,17 @@ from sglang.srt.managers.schedule_batch import (
     MultimodalInputs,
     flatten_nested_list,
 )
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.forward_batch_info import (
+    ForwardBatch,
+    ForwardMode,
+    PPProxyTensors,
+)
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
 from sglang.srt.models.gemma4_audio import Gemma4AudioEncoder
-from sglang.srt.models.gemma4_causal import Gemma4TextModel
+from sglang.srt.models.gemma4_causal import Gemma4TextModel, pp_filter_load_weight
 from sglang.srt.models.gemma4_vision import Gemma4VisionEncoder
 from sglang.srt.utils import add_prefix
 from sglang.srt.utils.hf_transformers_utils import get_processor
@@ -170,38 +177,46 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
         prefix: str = "",
     ) -> None:
         super().__init__(config=config)
+        self.pp_group = get_pp_group()
         self.config = config
         self.quant_config = quant_config
 
+        text_config = config.text_config
+
         prefix = add_prefix("model", prefix)
 
-        self.vision_tower = Gemma4VisionEncoder(
+        # Vision/audio encoders + their projection embedders are only consumed
-            config=config.vision_config,
+        # at the input-embedding stage, so they live on the first PP rank only.
-            quant_config=quant_config,
+        if self.pp_group.is_first_rank:
-            prefix=add_prefix("vision_tower", prefix),
+            self.vision_tower = Gemma4VisionEncoder(
-        )
+                config=config.vision_config,
-
-        self.embed_vision = Gemma4MultimodalEmbedder(
-            config.vision_config,
-            config.text_config,
-            quant_config=quant_config,
-            prefix=add_prefix("embed_vision", prefix),
-        )
-
-        # Audio components
-        if getattr(config, "audio_config", None) is not None:
-            self.audio_tower = Gemma4AudioEncoder(
-                config=config.audio_config,
                 quant_config=quant_config,
-                prefix=add_prefix("audio_tower", prefix),
+                prefix=add_prefix("vision_tower", prefix),
             )
-            self.embed_audio = Gemma4MultimodalEmbedder(
+            self.embed_vision = Gemma4MultimodalEmbedder(
-                config.audio_config,
+                config.vision_config,
                 config.text_config,
                 quant_config=quant_config,
-                prefix=add_prefix("embed_audio", prefix),
+                prefix=add_prefix("embed_vision", prefix),
             )
+            if getattr(config, "audio_config", None) is not None:
+                self.audio_tower = Gemma4AudioEncoder(
+                    config=config.audio_config,
+                    quant_config=quant_config,
+                    prefix=add_prefix("audio_tower", prefix),
+                )
+                self.embed_audio = Gemma4MultimodalEmbedder(
+                    config.audio_config,
+                    config.text_config,
+                    quant_config=quant_config,
+                    prefix=add_prefix("embed_audio", prefix),
+                )
+            else:
+                self.audio_tower = None
+                self.embed_audio = None
         else:
+            self.vision_tower = PPMissingLayer()
+            self.embed_vision = PPMissingLayer()
             self.audio_tower = None
             self.embed_audio = None
 
@@ -212,13 +227,31 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
             config.text_config.vocab_size,
         )
 
-        # Text model
+        # Text model — internal Gemma4TextModel is already PP-aware.
         self.language_model = Gemma4TextModel(
             config.text_config,
             quant_config,
             prefix=add_prefix("language_model", prefix),
         )
 
+        # Tied embeddings: under PP the embed_tokens lives on the first rank
+        # while logits run on the last rank, so we can't reuse the embedding
+        # module directly.  For PP=1 keep the original tying; for PP>1
+        # materialize a real ParallelLMHead on the last rank and route the
+        # checkpoint embedding into it during load_weights.
+        text_tie = getattr(text_config, "tie_word_embeddings", True)
+        if self.pp_group.world_size == 1 and text_tie:
+            self.lm_head = self.language_model.embed_tokens
+        elif self.pp_group.is_last_rank:
+            self.lm_head = ParallelLMHead(
+                text_config.vocab_size,
+                text_config.hidden_size,
+                quant_config=quant_config,
+                prefix=add_prefix("lm_head", prefix),
+            )
+        else:
+            self.lm_head = PPMissingLayer()
+
         # Create logits processor for the multimodal model
         self.logits_processor = LogitsProcessor(config.text_config)
         self.capture_aux_hidden_states = False
@@ -548,17 +581,26 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
         positions: torch.Tensor,
         forward_batch: ForwardBatch,
         input_embeds: torch.Tensor = None,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
         **kwargs: object,
-    ) -> LogitsProcessor:
+    ) -> Union[LogitsProcessor, PPProxyTensors]:
         """Forward pass for multimodal Gemma4."""
-        if (input_ids is None) ^ (input_embeds is not None):
+        is_first_rank = self.pp_group.is_first_rank
+        is_last_rank = self.pp_group.is_last_rank
+
+        # Only the first PP rank consumes input_ids/input_embeds; later stages
+        # receive activations through pp_proxy_tensors.
+        if is_first_rank and (input_ids is None) ^ (input_embeds is not None):
             raise ValueError(
                 "You must specify exactly one of input_ids or inputs_embeds"
             )
 
         positions += 1
         per_layer_inputs = None
-        if input_ids is not None:
+        # PLE table and the per-layer projection live on the first rank only,
+        # so non-first ranks must skip this and pull per_layer_inputs from the
+        # PP proxy (forwarded by Gemma4TextModel).
+        if is_first_rank and input_ids is not None:
             ple_ids = input_ids.clone()
             pad_id = self.config.text_config.pad_token_id
             ple_ids[input_ids == self.config.image_token_id] = pad_id
@@ -567,9 +609,8 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
             per_layer_inputs = self.get_per_layer_inputs(ple_ids)
 
         # Prepare bidirectional attention masks for image tokens during prefill.
-        # Gemma 4 uses bidirectional attention for image soft tokens.
+        # mm_inputs is preserved on every PP rank up to the first-rank embed
-        # Only TritonAttnBackend supports this; incompatible with CUDA Graph and
+        # routine, so each rank's attn_backend can install the mask locally.
-        # chunked prefill.
         if (
             forward_batch.forward_mode == ForwardMode.EXTEND
             and forward_batch.contains_image_inputs()
@@ -580,7 +621,8 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
                 mask_dtype=torch.bool,
             )
 
-        # Use general_mm_embed_routine for handling multimodal data
+        # general_mm_embed_routine already handles PP: it skips the embedding
+        # work on non-first ranks and forwards pp_proxy_tensors via **kwargs.
         hidden_states = general_mm_embed_routine(
             input_ids=input_ids,
             forward_batch=forward_batch,
@@ -592,24 +634,43 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
             },
             positions=positions,
             per_layer_inputs=per_layer_inputs,
+            pp_proxy_tensors=pp_proxy_tensors,
             **kwargs,
         )
 
+        if not is_last_rank:
+            # `hidden_states` is actually a PPProxyTensors flowing to the next
+            # stage; logits processing happens on the last rank only.
+            return hidden_states
+
         # Unpack aux_hidden_states if Eagle3 capture is active
         aux_hidden_states = None
         if self.capture_aux_hidden_states:
             hidden_states, aux_hidden_states = hidden_states
 
-        # Process hidden states through logits processor
+        # PP=1 keeps the original tied-weight behavior of using embed_tokens
+        # directly; under PP we route through the dedicated lm_head module.
+        head = (
+            self.language_model.embed_tokens
+            if self.pp_group.world_size == 1
+            and getattr(self.config.text_config, "tie_word_embeddings", True)
+            else self.lm_head
+        )
         return self.logits_processor(
             input_ids,
             hidden_states,
-            self.language_model.embed_tokens,
+            head,
             forward_batch,
             aux_hidden_states,
         )
 
     def tie_weights(self, recompute_mapping=False):
+        # Under PP, embed_tokens (first rank) and lm_head (last rank) live on
+        # different processes, so HF's automatic tying would crash on the
+        # PPMissingLayer side.  load_weights routes the embedding into lm_head
+        # on the last rank explicitly, so the tie is a no-op under PP.
+        if self.pp_group.world_size > 1:
+            return
         return self.language_model.tie_weights()
 
     # Standard stacked-params mapping for fused QKV / GateUp linears
@@ -764,6 +825,10 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
                 full = f"{mod_name}.{buf_name}" if mod_name else buf_name
                 non_persistent_buffers.add(full)
 
+        text_tie = getattr(self.config.text_config, "tie_word_embeddings", True)
+        start_layer = self.language_model.start_layer
+        end_layer = self.language_model.end_layer
+
         loaded_params: Set[str] = set()
 
         for name, loaded_weight in weights:
@@ -776,6 +841,29 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
 
             name = re.sub(r"^model\.", "", name)
 
+            if pp_filter_load_weight(
+                name,
+                loaded_weight,
+                pp_group=self.pp_group,
+                start_layer=start_layer,
+                end_layer=end_layer,
+                params_dict=params_dict,
+                loaded_params=loaded_params,
+                tie_word_embeddings=text_tie,
+                embed_weight_name="language_model.embed_tokens.weight",
+                first_rank_only_patterns=(
+                    "language_model.embed_tokens",
+                    "language_model.per_layer_model_projection",
+                    "language_model.per_layer_projection_norm",
+                    "vision_tower.",
+                    "embed_vision.",
+                    "audio_tower.",
+                    "embed_audio.",
+                ),
+                last_rank_only_prefixes=("language_model.norm.", "lm_head."),
+            ):
+                continue
+
             # HF has router.per_expert_scale and experts.* on the decoder layer;
             # remap into our moe.* subtree since Gemma4MoE owns both.
             name = name.replace(".router.per_expert_scale", ".moe.per_expert_scale")
@@ -948,6 +1036,18 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
         return self.language_model.embed_tokens.weight
 
     def get_embed_and_head(self):
+        if self.pp_group.world_size > 1:
+            # Under PP, embed_tokens lives on the first rank and lm_head on the
+            # last; neither rank holds both tensors, so we can't return the
+            # pair locally without a cross-stage gather.  Callers (RL weight
+            # sync, remote weight loader) currently assume a single-rank view —
+            # fail loudly rather than dereference a PPMissingLayer.
+            raise NotImplementedError(
+                "get_embed_and_head() is not implemented for Gemma4 "
+                "multimodal under pipeline parallelism. embed_tokens lives "
+                "on the first PP rank and lm_head on the last; use "
+                "--pp-size 1 if you need this API."
+            )
         embed = self.language_model.embed_tokens.weight
         # Gemma4 ties word embeddings, so embed_tokens serves as lm_head
         return embed, embed

python/sglang/srt/server_args.py

@@ -2194,7 +2194,10 @@ class ServerArgs:
                 f"Disable hybrid SWA memory for {model_arch} as it is not yet supported."
             )
             self.disable_hybrid_swa_memory = True
-        elif model_arch == "Gemma4ForConditionalGeneration":
+        elif model_arch in (
+            "Gemma4ForConditionalGeneration",
+            "Gemma4ForCausalLM",
+        ):
             default_attention_backend = (
                 "trtllm_mha" if is_sm100_supported() else "triton"
             )

python/sglang/test/test_utils.py

@@ -69,6 +69,8 @@ DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST = "Qwen/Qwen3-Next-80B-A3B-Instruct"
 # 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"
+DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP = "google/gemma-4-26B-A4B-it"
+DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP = "google/gemma-4-E4B-it"
 
 # NVFP4 models
 DEFAULT_DEEPSEEK_NVFP4_MODEL_FOR_TEST = "nvidia/DeepSeek-V3-0324-FP4"

test/registered/distributed/test_pp_single_node.py

@@ -4,6 +4,9 @@ python3 -m unittest test_pp_single_node.TestPPAccuracy.test_gsm8k
 python3 -m unittest test_pp_single_node.TestQwenPPAccuracy.test_pp_consistency
 python3 -m unittest test_pp_single_node.TestFixedBugs.test_chunked_prefill_with_small_bs
 python3 -m unittest test_pp_single_node.TestQwenVLPPAccuracy.test_mmmu
+python3 -m unittest test_pp_single_node.TestGemma4PPAccuracy.test_gsm8k
+python3 -m unittest test_pp_single_node.TestGemma4PPAccuracy.test_mmmu
+python3 -m unittest test_pp_single_node.TestGemma4PLEPPAccuracy.test_gsm8k
 python3 -m unittest test_pp_single_node.TestPPMixedChunk.test_gsm8k
 """
 
@@ -21,6 +24,8 @@ from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
     DEFAULT_MLA_MODEL_NAME_FOR_TEST,
     DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP,
+    DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP,
     DEFAULT_MODEL_NAME_FOR_TEST_GLM_41V_PP,
     DEFAULT_MODEL_NAME_FOR_TEST_VL_PP,
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -202,6 +207,137 @@ class TestQwenVLPPAccuracy(unittest.TestCase):
         self.assertGreater(metrics["score"], 0.26)
 
 
+@unittest.skipIf(
+    is_in_amd_ci(),
+    "Gemma4 PP not yet validated on AMD",
+)
+class TestGemma4PPAccuracy(unittest.TestCase):
+    """End-to-end PP=2 accuracy gate for Gemma4 multimodal.
+
+    Gemma4 has full-attention layers with head_dim=512 (FA's max is 256), so
+    sglang auto-selects the triton attention backend; no manual flag needed.
+    The 26B BF16 model splits to ~26 GB per stage under PP=2, well within an
+    H100's 80 GB.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP
+        cls.base_url = "http://127.0.0.1:23333"
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--tp-size",
+                1,
+                "--pp-size",
+                2,
+                "--trust-remote-code",
+                "--enable-multimodal",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        # Gemma4 is instruction-tuned and doesn't follow few-shot completion
+        # prompts well — use the chat API (default in run_eval), which scores
+        # ~0.98 on this model vs ~0.44 with api="completion".
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            num_examples=200,
+            num_threads=32,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        # Chat-API baseline ~0.98; gate well below to absorb sample-noise
+        # without missing a real PP-routing regression (pre-PP-fix the model
+        # produced garbage outputs scoring ≈ 0).
+        self.assertGreaterEqual(metrics["score"], 0.90)
+        # Wait a little bit so that the memory check happens.
+        time.sleep(4)
+
+    @unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
+    def test_mmmu(self):
+        # Multimodal accuracy gate covering the vision_tower → embed_vision
+        # (first rank) → PP-proxy handoff → LM tail (last rank) chain.
+        # Measured 0.71 on 200 examples; full eval (~900 questions) takes
+        # ~5-7 min on H100 so this is manual-only.
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mmmu",
+            num_examples=None,
+            num_threads=32,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        # Measured 0.72 on this setup; published Gemma-4-26B MMMU lies in
+        # 0.69-0.73.  Gate 0.65 leaves ~5 SE of headroom (SE on 900 binary
+        # samples ≈ 0.015) while still catching mid-grade vision/PP
+        # regressions, not just complete breakage.
+        self.assertGreater(metrics["score"], 0.65)
+
+
+@unittest.skipIf(
+    is_in_amd_ci(),
+    "Gemma4 PP not yet validated on AMD",
+)
+class TestGemma4PLEPPAccuracy(unittest.TestCase):
+    """PP=2 coverage for Gemma4 PLE variants (per_layer_inputs proxy path).
+
+    26B-A4B has ``hidden_size_per_layer_input=0`` so the default Gemma4 PP
+    test never crosses the PLE branch.  Cuda graph + PLE corrupts outputs
+    (the runner's hardcoded ``{hidden_states, residual}`` PP-proxy schema
+    drops ``per_layer_inputs``), so this test pins the eager configuration.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP
+        cls.base_url = "http://127.0.0.1:23339"
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--tp-size",
+                1,
+                "--pp-size",
+                2,
+                "--trust-remote-code",
+                "--enable-multimodal",
+                # Required for PLE under PP — see Gemma4TextModel guard.
+                "--disable-cuda-graph",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        # Eager-path baseline ~0.92; gate 0.80 catches PLE breakage
+        # (corruption collapses score to ~0).
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            num_examples=100,
+            num_threads=32,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["score"], 0.80)
+        time.sleep(4)
+
+
 class TestQwenPPAccuracy(unittest.TestCase):
     @classmethod
     def setUpClass(cls):