[diffusion] chore: reuse SRT SigLIP in Pi0.5 (#34992)

docs/cookbook/vla/OpenPI/Pi0.5.mdx

@@ -18,6 +18,8 @@ Pi0.5 is an OpenPI / LeRobot diffusion Vision-Language-Action (dVLA) policy. It
 
 SGLang serves Pi0.5 through the native `multimodal_gen` runtime. The implementation uses a SigLIP/PaliGemma prefix encoder and a Gemma action expert: the prefix is encoded once, then the action expert runs the flow-matching denoising loop. This is not a token decode workload, so the Pi0.5 path does not use the LLM sampler, logits processor, token streaming, paged decode KV cache, or a separate SRT serving engine.
 
+The SigLIP vision tower, PaliGemma language stack, and action expert are all SGLang-native modules. Transformers is used for checkpoint configuration and tokenization, not for the runtime neural network.
+
 The prefix encoder covers both stages of observation encoding: SigLIP turns resized camera pixels into continuous patch embeddings, then the PaliGemma transformer jointly encodes those patches with tokenized task/state inputs and produces per-layer prefix K/V. At flow timestep `t`, the action expert projects the noisy continuous action chunk `x_t` into action embeddings. Its queries attend to both the fixed prefix K/V and the current action K/V, while the timestep follows a separate sinusoidal-MLP path and conditions every action-expert layer through AdaRMSNorm gates.
 
 Supported public checkpoints:

python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import math
 from dataclasses import dataclass
+from functools import partial
 from typing import Any, Literal
 
 import torch
@@ -12,8 +13,7 @@ import torch.nn.functional as F
 from torch import Tensor, nn
 from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.models.auto import CONFIG_MAPPING
-from transformers.models.gemma.modeling_gemma import GemmaConfig
+from transformers.models.gemma.configuration_gemma import GemmaConfig
-from transformers.models.paligemma.modeling_paligemma import PaliGemmaModel
 
 from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
 from sglang.multimodal_gen.runtime.distributed.parallel_state import (
@@ -22,6 +22,7 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_ulysses_parallel_world_size,
     model_parallel_is_initialized,
 )
+from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
 from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention
 from sglang.multimodal_gen.runtime.layers.linear import (
     MergedColumnParallelLinear,
@@ -30,12 +31,16 @@ from sglang.multimodal_gen.runtime.layers.linear import (
 )
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import RotaryEmbedding
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 from sglang.multimodal_gen.runtime.vla.prefix_cache import VLADensePrefixCache
 from sglang.srt.layers.activation import GeluAndMul
 from sglang.srt.layers.rotary_embedding import (
     apply_rotary_pos_emb as native_apply_rotary_pos_emb,
 )
+from sglang.srt.models.siglip import SiglipVisionModel
 
 
 def config_compute_dtype(config: GemmaConfig) -> torch.dtype | None:
@@ -106,73 +111,6 @@ def _use_ulysses_action_attention(num_heads: int) -> bool:
     )
 
 
-class Pi05SiglipAttention(nn.Module):
-    def __init__(self, attention: nn.Module):
-        super().__init__()
-        self.embed_dim = attention.embed_dim
-        self.num_heads = attention.num_heads
-        self.head_dim = attention.head_dim
-        self.scale = getattr(attention, "scale", self.head_dim**-0.5)
-        self.dropout = getattr(attention, "dropout", 0.0)
-        self.q_proj = attention.q_proj
-        self.k_proj = attention.k_proj
-        self.v_proj = attention.v_proj
-        self.out_proj = attention.out_proj
-        self.attn = LocalAttention(
-            num_heads=self.num_heads,
-            head_size=self.head_dim,
-            num_kv_heads=self.num_heads,
-            softmax_scale=self.scale,
-            causal=False,
-            supported_attention_backends={
-                AttentionBackendEnum.FA,
-                AttentionBackendEnum.FA2,
-                AttentionBackendEnum.TORCH_SDPA,
-            },
-            compute_dtype=self.q_proj.weight.dtype,
-            allow_cudnn_sdp=True,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        output_attentions: bool = False,
-        **kwargs,
-    ) -> tuple[torch.Tensor, None]:
-        input_shape = hidden_states.shape[:-1]
-        query_states = self.q_proj(hidden_states).view(
-            *input_shape,
-            self.num_heads,
-            self.head_dim,
-        )
-        key_states = self.k_proj(hidden_states).view(
-            *input_shape,
-            self.num_heads,
-            self.head_dim,
-        )
-        value_states = self.v_proj(hidden_states).view(
-            *input_shape,
-            self.num_heads,
-            self.head_dim,
-        )
-        attn_output = self.attn(
-            query_states,
-            key_states,
-            value_states,
-            attn_mask=attention_mask,
-        )
-        attn_output = attn_output.reshape(*input_shape, self.embed_dim).contiguous()
-        return self.out_proj(attn_output), None
-
-
-def patch_siglip_vision_attention_to_native(vision_model: nn.Module) -> None:
-    for layer in vision_model.encoder.layers:
-        if isinstance(layer.self_attn, Pi05SiglipAttention):
-            continue
-        layer.self_attn = Pi05SiglipAttention(layer.self_attn)
-
-
 class PiGemmaRMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6, cond_dim: int | None = None):
         super().__init__()
@@ -788,15 +726,51 @@ class PiGemmaForCausalLM(nn.Module):
         self.lm_head = None
 
 
-class PaliGemmaModelWithPiGemma(PaliGemmaModel):
+class PaliGemmaMultiModalProjector(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.linear = nn.Linear(
+            config.vision_config.hidden_size,
+            config.vision_config.projection_dim,
+            bias=True,
+        )
+
+    def forward(self, image_features: torch.Tensor) -> torch.Tensor:
+        return self.linear(image_features)
+
+
+class Pi05SiglipVisionModel(SiglipVisionModel, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["vision_model.encoder.layers"]
+
+
+class PaliGemmaModelWithPiGemma(nn.Module):
     def __init__(self, config, *, tensor_parallel: bool = False):
-        super().__init__(config)
+        super().__init__()
-        del self.language_model
+        self.config = config
+        self.vision_tower = Pi05SiglipVisionModel(
+            config.vision_config,
+            act_layer=partial(get_act_fn, config.vision_config.hidden_act),
+            qkv_backend="sdpa",
+            flatten_batch=False,
+            use_data_parallel=True,
+        )
+        self.multi_modal_projector = PaliGemmaMultiModalProjector(config)
         self.language_model = PiGemmaModel(
             config.text_config,
             tensor_parallel=tensor_parallel,
         )
 
+    def get_image_features(
+        self, pixel_values: torch.Tensor
+    ) -> BaseModelOutputWithPooling:
+        vision_features = self.vision_tower(pixel_values)
+        image_features = self.multi_modal_projector(vision_features)
+        return BaseModelOutputWithPooling(
+            last_hidden_state=vision_features,
+            pooler_output=image_features,
+        )
+
 
 class PaliGemmaForConditionalGenerationWithPiGemma(nn.Module):
     def __init__(self, config, *, tensor_parallel: bool = False):
@@ -907,34 +881,6 @@ def prepare_optional_full_attention_mask(
     return torch.where(masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
 
 
-def siglip_vision_forward_with_openpi_dtype(
-    self,
-    pixel_values,
-    interpolate_pos_encoding: bool | None = False,
-    **kwargs,
-) -> BaseModelOutputWithPooling:
-    hidden_states = self.embeddings(
-        pixel_values,
-        interpolate_pos_encoding=interpolate_pos_encoding,
-    )
-    if (
-        len(self.encoder.layers) > 0
-        and self.encoder.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16
-    ):
-        hidden_states = hidden_states.to(torch.bfloat16)
-
-    encoder_outputs = self.encoder(inputs_embeds=hidden_states, **kwargs)
-    last_hidden_state = encoder_outputs.last_hidden_state
-    last_hidden_state = self.post_layernorm(last_hidden_state)
-    pooler_output = self.head(last_hidden_state) if self.use_head else None
-    return BaseModelOutputWithPooling(
-        last_hidden_state=last_hidden_state,
-        pooler_output=pooler_output,
-        hidden_states=encoder_outputs.hidden_states,
-        attentions=encoder_outputs.attentions,
-    )
-
-
 def compute_layer_complete(
     inputs_embeds,
     attention_mask,
@@ -1059,12 +1005,6 @@ class PaliGemmaWithExpertModel(nn.Module):
                 config=vlm_config_hf,
                 tensor_parallel=prefix_tensor_parallel,
             )
-            vision_tower = self.paligemma.model.vision_tower
-            vision_model = getattr(vision_tower, "vision_model", vision_tower)
-            vision_model.forward = siglip_vision_forward_with_openpi_dtype.__get__(
-                vision_model,
-                type(vision_model),
-            )
             self.paligemma.lm_head = None
 
         if runtime_role in ("all", "action"):
@@ -1090,14 +1030,6 @@ class PaliGemmaWithExpertModel(nn.Module):
             self.gemma_expert.lm_head = None
             self.gemma_expert.model.embed_tokens = None
         self.to_selected_dtype(precision)
-        self.patch_native_attention_after_dtype_finalize()
-
-    def patch_native_attention_after_dtype_finalize(self) -> None:
-        if self.paligemma is None:
-            return
-        vision_tower = self.paligemma.model.vision_tower
-        vision_model = getattr(vision_tower, "vision_model", vision_tower)
-        patch_siglip_vision_attention_to_native(vision_model)
 
     def to_selected_dtype(
         self, precision: Literal["bfloat16", "float32"] = "bfloat16"

python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py

@@ -569,6 +569,7 @@ class Pi05PolicyModel(nn.Module):
         candidates = [key]
         replacements = {
             ".vision_tower.vision_model.": ".vision_tower.",
+            ".self_attn.out_proj.": ".self_attn.proj.",
             ".paligemma.language_model.": ".paligemma.model.language_model.",
             ".paligemma.vision_tower.": ".paligemma.model.vision_tower.",
             ".paligemma.multi_modal_projector.": (
@@ -576,8 +577,9 @@ class Pi05PolicyModel(nn.Module):
             ),
         }
         for old, new in replacements.items():
-            if old in key:
+            for candidate in list(candidates):
-                candidates.append(key.replace(old, new))
+                if old in candidate:
+                    candidates.append(candidate.replace(old, new))
 
         if key in {
             "paligemma_with_expert.paligemma.lm_head.weight",

python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py

@@ -10,8 +10,7 @@ import sglang.multimodal_gen.runtime.models.vlas.pi05_policy as pi05_policy_modu
 from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
 from sglang.multimodal_gen.runtime.models.vlas.pi05_core import (
     Pi05CoreModel,
-    Pi05SiglipAttention,
+    Pi05SiglipVisionModel,
-    patch_siglip_vision_attention_to_native,
 )
 from sglang.multimodal_gen.runtime.models.vlas.pi05_policy import (
     Pi05CheckpointManifest,
@@ -30,6 +29,8 @@ from sglang.multimodal_gen.runtime.vla.prefix_cache import (
     PrefixContext,
     VLADensePrefixCache,
 )
+from sglang.srt.models.siglip import SiglipVisionModel
+from sglang.srt.runtime_context import get_context
 
 
 def _prefix_context(value: float, digest: str | None) -> PrefixContext:
@@ -183,30 +184,64 @@ def test_action_parallel_info_reports_single_rank_without_process_group():
     }
 
 
-class _FakeSiglipAttention(nn.Module):
+def test_pi05_siglip_reuses_srt_model_with_layerwise_groups():
-    def __init__(self):
+    config = SimpleNamespace(
-        super().__init__()
+        hidden_size=8,
-        self.embed_dim = 8
+        intermediate_size=16,
-        self.num_heads = 2
+        num_hidden_layers=1,
-        self.head_dim = 4
+        num_attention_heads=2,
-        self.scale = self.head_dim**-0.5
+        layer_norm_eps=1e-6,
-        self.dropout = 0.0
+        image_size=4,
-        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        patch_size=2,
-        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        num_channels=3,
-        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        hidden_act="gelu_pytorch_tanh",
-        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+    )
+    with get_context().override_server_args():
+        model = Pi05SiglipVisionModel(
+            config,
+            act_layer=lambda: nn.GELU(approximate="tanh"),
+            qkv_backend="sdpa",
+            flatten_batch=False,
+            use_data_parallel=True,
+        )
+
+    state_keys = set(model.state_dict())
+    prefix = "vision_model.encoder.layers.0.self_attn"
+    vision_model = model.vision_model
+    layer = vision_model.encoder.layers[0]
+
+    assert isinstance(model, SiglipVisionModel)
+    assert f"{prefix}.qkv_proj.weight" in state_keys
+    assert f"{prefix}.proj.weight" in state_keys
+    assert vision_model.embeddings.position_embedding.tp_size == 1
+    assert layer.self_attn.tp_size == 1
+    assert layer.self_attn.qkv_backend.flatten_batch is False
+    assert layer.mlp.fc1.tp_size == 1
+    assert layer.mlp.fc2.tp_size == 1
+    assert isinstance(layer.mlp.act, nn.GELU)
+    assert layer.mlp.act.approximate == "tanh"
+    assert model.device == vision_model.embeddings.patch_embedding.weight.device
+    assert model.layer_names == ["vision_model.encoder.layers"]
 
 
-def test_siglip_attention_patch_uses_native_wrapper_once():
+def test_pi05_siglip_checkpoint_names_map_to_srt_layers():
-    layer = SimpleNamespace(self_attn=_FakeSiglipAttention())
+    source_prefix = (
-    vision_model = SimpleNamespace(encoder=SimpleNamespace(layers=[layer]))
+        "paligemma_with_expert.paligemma.vision_tower.vision_model."
+        "encoder.layers.0.self_attn"
+    )
+    target_prefix = (
+        "paligemma_with_expert.paligemma.model.vision_tower.vision_model."
+        "encoder.layers.0.self_attn"
+    )
 
-    patch_siglip_vision_attention_to_native(vision_model)
+    assert (
-    first = layer.self_attn
+        f"{target_prefix}.qkv_proj.weight",
-    patch_siglip_vision_attention_to_native(vision_model)
+        "q",
-
+    ) in Pi05PolicyModel._candidate_target_weights(f"{source_prefix}.q_proj.weight")
-    assert isinstance(first, Pi05SiglipAttention)
+    assert (
-    assert layer.self_attn is first
+        f"{target_prefix}.proj.weight",
+        None,
+    ) in Pi05PolicyModel._candidate_target_weights(f"{source_prefix}.out_proj.weight")
 
 
 def test_prefix_language_embedding_matches_openpi_scale():

python/sglang/srt/models/siglip.py

@@ -2,7 +2,7 @@
 # https://github.com/huggingface/transformers/blob/af9b2eaa54c150741f298d6db939af6328e1dc38/src/transformers/models/siglip/modeling_siglip.py
 
 from functools import partial
-from typing import Optional, Type, Union
+from typing import Callable, Optional, Type, Union
 
 import torch
 import torch.nn as nn
@@ -14,13 +14,14 @@ from sglang.srt.layers.conv import Conv2dLayer
 from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix
 
 
 # Adapted from transformers.models.siglip.modeling_siglip.SiglipVisionTransformer
 class SiglipVisionEmbeddings(nn.Module):
 
-    def __init__(self, config: SiglipVisionConfig):
+    def __init__(self, config: SiglipVisionConfig, use_data_parallel: bool = False):
         super().__init__()
         self.config = config
         self.embed_dim = config.hidden_size
@@ -38,7 +39,9 @@ class SiglipVisionEmbeddings(nn.Module):
         self.num_patches = (self.image_size // self.patch_size) ** 2
         self.num_positions = self.num_patches
         self.position_embedding = VocabParallelEmbedding(
-            self.num_positions, self.embed_dim
+            self.num_positions,
+            self.embed_dim,
+            enable_tp=not use_data_parallel,
         )
         self.register_buffer(
             "position_ids",
@@ -64,16 +67,21 @@ class SiglipMLP(nn.Module):
     def __init__(
         self,
         config,
-        act_layer: Type[nn.Module] = QuickGELU,
+        act_layer: Callable[[], nn.Module] = QuickGELU,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
+        use_data_parallel: bool = False,
     ):
         super().__init__()
+        tp_size = 1 if use_data_parallel else get_parallel().tp_size
+        tp_rank = 0 if use_data_parallel else get_parallel().tp_rank
         self.fc1 = ColumnParallelLinear(
             config.hidden_size,
             config.intermediate_size,
             quant_config=quant_config,
             prefix=add_prefix("fc1", prefix),
+            tp_size=tp_size,
+            tp_rank=tp_rank,
         )
         self.act = act_layer()
         self.fc2 = RowParallelLinear(
@@ -81,6 +89,8 @@ class SiglipMLP(nn.Module):
             config.hidden_size,
             quant_config=quant_config,
             prefix=add_prefix("fc2", prefix),
+            tp_size=tp_size,
+            tp_rank=tp_rank,
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -96,11 +106,13 @@ class SiglipEncoderLayer(nn.Module):
     def __init__(
         self,
         config: SiglipVisionConfig,
-        act_layer: Type[nn.Module] = QuickGELU,
+        act_layer: Callable[[], nn.Module] = QuickGELU,
         norm_layer: Type[nn.Module] = None,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
         qkv_backend: Optional[str] = None,
+        flatten_batch: bool = True,
+        use_data_parallel: bool = False,
     ) -> None:
         super().__init__()
         if norm_layer is None:
@@ -112,14 +124,16 @@ class SiglipEncoderLayer(nn.Module):
             num_heads=config.num_attention_heads,
             projection_size=config.hidden_size,
             use_qkv_parallel=True,
-            flatten_batch=True,
+            flatten_batch=flatten_batch,
             qkv_backend=qkv_backend,
+            use_data_parallel=use_data_parallel,
             quant_config=quant_config,
             prefix=add_prefix("self_attn", prefix),
         )
         self.mlp = SiglipMLP(
             config,
             act_layer=act_layer,
+            use_data_parallel=use_data_parallel,
             quant_config=quant_config,
             prefix=add_prefix("mlp", prefix),
         )
@@ -170,6 +184,9 @@ class SiglipEncoder(nn.Module):
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
         qkv_backend: Optional[str] = None,
+        act_layer: Callable[[], nn.Module] = QuickGELU,
+        flatten_batch: bool = True,
+        use_data_parallel: bool = False,
     ) -> None:
         super().__init__()
 
@@ -183,6 +200,9 @@ class SiglipEncoder(nn.Module):
                     config=config,
                     norm_layer=norm_layer,
                     qkv_backend=qkv_backend,
+                    act_layer=act_layer,
+                    flatten_batch=flatten_batch,
+                    use_data_parallel=use_data_parallel,
                     quant_config=quant_config,
                     prefix=add_prefix(f"layers.{layer_idx}", prefix),
                 )
@@ -220,17 +240,25 @@ class SiglipVisionTransformer(nn.Module):
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
         qkv_backend: Optional[str] = None,
+        act_layer: Callable[[], nn.Module] = QuickGELU,
+        flatten_batch: bool = True,
+        use_data_parallel: bool = False,
     ) -> None:
         super().__init__()
 
         self.config = config
         embed_dim = config.hidden_size
 
-        self.embeddings = SiglipVisionEmbeddings(config)
+        self.embeddings = SiglipVisionEmbeddings(
+            config, use_data_parallel=use_data_parallel
+        )
 
         self.encoder = SiglipEncoder(
             config=config,
             qkv_backend=qkv_backend,
+            act_layer=act_layer,
+            flatten_batch=flatten_batch,
+            use_data_parallel=use_data_parallel,
             quant_config=quant_config,
             prefix=add_prefix("encoder", prefix),
         )
@@ -247,13 +275,15 @@ class SiglipVisionTransformer(nn.Module):
 
     @property
     def device(self) -> torch.device:
-        return self.encoder.layers[0].layer_norm1.weight.device
+        return self.embeddings.patch_embedding.weight.device
 
     def forward(
         self,
         pixel_values: torch.Tensor,
     ) -> torch.Tensor:
-        hidden_states = self.embeddings(pixel_values.to(self.device))
+        hidden_states = self.embeddings(pixel_values.to(self.device)).to(
+            self.post_layernorm.weight.dtype
+        )
 
         return_all_hidden_states = False
 
@@ -275,11 +305,17 @@ class SiglipVisionModel(nn.Module):
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
         qkv_backend: Optional[str] = None,
+        act_layer: Callable[[], nn.Module] = QuickGELU,
+        flatten_batch: bool = True,
+        use_data_parallel: bool = False,
     ):
         super().__init__()
         self.vision_model = SiglipVisionTransformer(
             config,
             qkv_backend=qkv_backend,
+            act_layer=act_layer,
+            flatten_batch=flatten_batch,
+            use_data_parallel=use_data_parallel,
             quant_config=quant_config,
             prefix=add_prefix("vision_model", prefix),
         )