model: support Pi0.5 (#30633)

python/pyproject.toml

@@ -104,6 +104,7 @@ diffusion = [
   "imageio==2.36.0",
   "imageio-ffmpeg==0.5.1",
   "moviepy>=2.0.0",
+  "msgpack",
   "nvidia-modelopt",
   "opencv-python-headless==4.10.0.84",
   "PyYAML==6.0.1",

python/sglang/cli/serve.py

@@ -46,13 +46,21 @@ def _extract_model_type_override(extra_argv):
     return model_type, filtered_argv
 
 
+def _normalize_positional_model_path(extra_argv):
+    """Allow `sglang serve <model>` while preserving existing flag parsing."""
+    if extra_argv and not extra_argv[0].startswith("-"):
+        return ["--model-path", extra_argv[0], *extra_argv[1:]], True
+    return extra_argv, False
+
+
 def serve(args, extra_argv):
     if any(h in extra_argv for h in ("-h", "--help")):
         # Since the server type is determined by the model, and we don't have a model path,
         # we can't show the exact help. Instead, we show a general help message and then
         # the help for both possible server types.
         print(
-            "Usage: sglang serve --model-path <model-name-or-path> [additional-arguments]\n\n"
+            "Usage: sglang serve <model-name-or-path> [additional-arguments]\n"
+            "   or: sglang serve --model-path <model-name-or-path> [additional-arguments]\n\n"
             "This command can launch either a standard language model server or a diffusion model server.\n"
             "The server type is determined by the --model-path.\n"
             "Optional override: --model-type {auto,llm,diffusion} "
@@ -91,6 +99,9 @@ def serve(args, extra_argv):
     load_plugins()
 
     model_type, dispatch_argv = _extract_model_type_override(extra_argv)
+    dispatch_argv, positional_model_path = _normalize_positional_model_path(
+        dispatch_argv
+    )
     model_path = get_model_path(dispatch_argv)
     try:
         if model_type == "auto":
@@ -116,6 +127,8 @@ def serve(args, extra_argv):
             )
             add_multimodal_gen_serve_args(parser)
             parsed_args, remaining_argv = parser.parse_known_args(dispatch_argv)
+            if positional_model_path:
+                parsed_args._sglang_explicit_arg_names = {"model_path"}
 
             execute_serve_cmd(parsed_args, remaining_argv)
         else:

python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py

@@ -40,6 +40,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
     LTX23PipelineConfig,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import (
     StableDiffusion3PipelineConfig,
@@ -71,6 +72,7 @@ __all__ = [
     "SanaPipelineConfig",
     "SlidingTileAttnConfig",
     "MOVAPipelineConfig",
+    "Pi05PipelineConfig",
     "StableDiffusion3PipelineConfig",
     "WanT2V480PConfig",
     "WanI2V480PConfig",

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -59,6 +59,7 @@ class ModelTaskType(Enum):
     I2I = auto()  # Image to Image
     TI2I = auto()  # Image to Image or Text-Image to Image
     I2M = auto()  # Image to Mesh
+    VLA_ACTION = auto()  # Vision-language-action policy output
 
     def is_image_gen(self) -> bool:
         return (
@@ -67,6 +68,22 @@ class ModelTaskType(Enum):
             or self == ModelTaskType.TI2I
         )
 
+    def is_action_gen(self) -> bool:
+        return self == ModelTaskType.VLA_ACTION
+
+    def is_mesh_gen(self) -> bool:
+        return self == ModelTaskType.I2M
+
+    def is_video_gen(self) -> bool:
+        return (
+            self == ModelTaskType.I2V
+            or self == ModelTaskType.T2V
+            or self == ModelTaskType.TI2V
+        )
+
+    def is_visual_gen(self) -> bool:
+        return self.is_image_gen() or self.is_video_gen()
+
     def requires_image_input(self) -> bool:
         return (
             self == ModelTaskType.I2V
@@ -81,15 +98,17 @@ class ModelTaskType(Enum):
             or self == ModelTaskType.TI2I
             or self == ModelTaskType.TI2V
             or self == ModelTaskType.I2M
+            or self == ModelTaskType.VLA_ACTION
         )
 
     def data_type(self) -> DataType:
-        if self == ModelTaskType.I2M:
+        if self.is_action_gen():
+            return DataType.ACTION
+        if self.is_mesh_gen():
             return DataType.MESH
         if self.is_image_gen():
             return DataType.IMAGE
-        else:
-            return DataType.VIDEO
+        return DataType.VIDEO
 
 
 class STA_Mode(str, Enum):
@@ -374,6 +393,10 @@ class PipelineConfig:
         """
         return self.task_type in (ModelTaskType.T2I, ModelTaskType.T2V)
 
+    def supports_native_grouped_requests(self):
+        """Return whether dynamic batches should run as grouped Req lists."""
+        return False
+
     def estimate_request_cost(self, batch) -> float:
         """Return the relative cost used for batching admission caps.
 
@@ -888,6 +911,8 @@ class PipelineConfig:
         pipeline_config_or_path: str | PipelineConfig | dict[str, Any] | None = (
             kwargs.get(prefix_with_dot + "pipeline_config", None)
             or kwargs.get("pipeline_config")
+            or kwargs.get(prefix_with_dot + "pipeline_config_path", None)
+            or kwargs.get("pipeline_config_path")
         )
         if model_path is None:
             raise ValueError("model_path is required in kwargs")
@@ -943,36 +968,52 @@ class PipelineConfig:
                 model_id=kwargs.get("model_id"),
             )
             if model_info is None:
-                raise ValueError(
-                    f"Could not get model info for '{model_path}'. "
-                    f"If using a safetensors file, please specify pipeline_class_name"
-                )
-            # 1.5. Adjust pipeline config for fine-tuned VAE if needed
-            pipeline_config_cls = model_info.pipeline_config_cls
-            # If an explicit pipeline_class_name refines the model-default config
-            # (e.g. SanaWMRealtimePipeline -> SanaWMRealtimeConfig, a subclass of
-            # the model-resolved SanaWMPipelineConfig), prefer the pipeline's own
-            # config so realtime-only wiring (the /v1/realtime_video adapter) is
-            # selected. Only applies when the explicit config strictly subclasses
-            # the model default, so non-realtime pipelines are unaffected.
-            if pipeline_class_name:
-                explicit_config_classes = get_pipeline_config_classes(
-                    pipeline_class_name
-                )
-                if explicit_config_classes is not None:
-                    explicit_config_cls = explicit_config_classes[0]
-                    if (
-                        isinstance(explicit_config_cls, type)
-                        and isinstance(pipeline_config_cls, type)
-                        and explicit_config_cls is not pipeline_config_cls
-                        and issubclass(explicit_config_cls, pipeline_config_cls)
-                    ):
+                if pipeline_class_name:
+                    config_classes = get_pipeline_config_classes(pipeline_class_name)
+                    if config_classes is not None:
+                        pipeline_config_cls = config_classes[0]
                         logger.info(
-                            f"Refining pipeline config {pipeline_config_cls.__name__} "
-                            f"-> {explicit_config_cls.__name__} for explicit "
-                            f"pipeline_class_name={pipeline_class_name}"
+                            "Using %s from explicit pipeline_class_name=%s",
+                            pipeline_config_cls.__name__,
+                            pipeline_class_name,
                         )
-                        pipeline_config_cls = explicit_config_cls
+                    else:
+                        raise ValueError(
+                            f"Could not get model info for '{model_path}'. "
+                            "Please specify a valid model_id or pipeline_class_name."
+                        )
+                else:
+                    raise ValueError(
+                        f"Could not get model info for '{model_path}'. "
+                        f"If using a safetensors file, please specify pipeline_class_name"
+                    )
+            else:
+                # 1.5. Adjust pipeline config for fine-tuned VAE if needed
+                pipeline_config_cls = model_info.pipeline_config_cls
+                # If an explicit pipeline_class_name refines the model-default config
+                # (e.g. SanaWMRealtimePipeline -> SanaWMRealtimeConfig, a subclass of
+                # the model-resolved SanaWMPipelineConfig), prefer the pipeline's own
+                # config so realtime-only wiring (the /v1/realtime_video adapter) is
+                # selected. Only applies when the explicit config strictly subclasses
+                # the model default, so non-realtime pipelines are unaffected.
+                if pipeline_class_name:
+                    explicit_config_classes = get_pipeline_config_classes(
+                        pipeline_class_name
+                    )
+                    if explicit_config_classes is not None:
+                        explicit_config_cls = explicit_config_classes[0]
+                        if (
+                            isinstance(explicit_config_cls, type)
+                            and isinstance(pipeline_config_cls, type)
+                            and explicit_config_cls is not pipeline_config_cls
+                            and issubclass(explicit_config_cls, pipeline_config_cls)
+                        ):
+                            logger.info(
+                                f"Refining pipeline config {pipeline_config_cls.__name__} "
+                                f"-> {explicit_config_cls.__name__} for explicit "
+                                f"pipeline_class_name={pipeline_class_name}"
+                            )
+                            pipeline_config_cls = explicit_config_cls
         vae_path = kwargs.get(prefix_with_dot + "vae_path") or kwargs.get("vae_path")
         if vae_path is None:
             component_paths = kwargs.get(

python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.pipeline_configs.base import (
+    ModelTaskType,
+    PipelineConfig,
+)
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
+
+
+@dataclass
+class Pi05PipelineConfig(PipelineConfig):
+    """Configuration for OpenPI / LeRobot Pi0.5 action policies."""
+
+    task_type: ModelTaskType = ModelTaskType.VLA_ACTION
+    should_use_guidance: bool = False
+    enable_autocast: bool = True
+    generator_device: str | None = None
+
+    # OpenPI pi0.5 public checkpoint layout.
+    pi05: bool = True
+    paligemma_variant: str = "gemma_2b"
+    action_expert_variant: str = "gemma_300m"
+    max_token_len: int = 200
+    action_horizon: int = 50
+    action_dim: int = 32
+    state_dim: int = 32
+    output_action_dim: int = 32
+    n_action_steps: int = 50
+    default_num_inference_steps: int = 10
+    time_embedding_min_period: float = 4e-3
+    time_embedding_max_period: float = 4.0
+    tokenizer_name: str = "google/paligemma-3b-pt-224"
+
+    image_keys: tuple[str, ...] = (
+        "base_0_rgb",
+        "left_wrist_0_rgb",
+        "right_wrist_0_rgb",
+    )
+    empty_cameras: int = 0
+    image_size: tuple[int, int] = (224, 224)
+    image_normalization_mean: tuple[float, float, float] = (0.5, 0.5, 0.5)
+    image_normalization_std: tuple[float, float, float] = (0.5, 0.5, 0.5)
+
+    enable_global_prefix_cache: bool = False
+    enable_action_cuda_graph: bool = True
+    prefix_cache_max_entries: int = 1
+    prefix_cache_layout_version: str = "pi05-prefix-v1"
+    offload_prefix_image_encoder: bool = False
+    offload_prefix_image_encoder_after_embed: bool = False
+    offload_prefix_token_embedding: bool = False
+    offload_prefix_language_layers: bool = False
+    offload_prefix_language_layers_after_prefix: bool = False
+    offload_prefix_language_layer_count_after_prefix: int = 0
+    offload_prefix_language_layers_empty_cache: bool = True
+    offload_action_expert_after_denoise: bool = False
+    empty_cache_after_prefix: bool = False
+
+    # Prefix VLM and action expert are separate logical groups. The concrete
+    # process-group construction lands with the native model parallel kernels.
+    prefix_parallel_strategy: str = "tp"
+    action_parallel_strategy: str = "sp"
+    parallel_layout_version: str = "pi05-split-prefix-action-v1"
+
+    skip_unused_lm_head: bool = True
+    materialize_dtype: str = "bf16"
+    loader_component_map: dict[str, tuple[str, ...]] = field(
+        default_factory=lambda: {
+            "vision_tower": ("paligemma_with_expert.paligemma.model.vision_tower.",),
+            "paligemma": ("paligemma_with_expert.paligemma.model.language_model.",),
+            "multi_modal_projector": (
+                "paligemma_with_expert.paligemma.model.multi_modal_projector.",
+            ),
+            "action_expert": ("paligemma_with_expert.gemma_expert.",),
+            "action_heads": (
+                "action_in_proj.",
+                "action_out_proj.",
+                "time_mlp_in.",
+                "time_mlp_out.",
+            ),
+        }
+    )
+
+    def supports_dynamic_batching(self):
+        return True
+
+    def supports_native_grouped_requests(self):
+        return True
+
+    def estimate_request_cost(self, batch) -> float:
+        return float(
+            self.action_horizon * self.action_dim * self.default_num_inference_steps
+        )
+
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig()

python/sglang/multimodal_gen/configs/sample/__init__.py

@@ -4,10 +4,14 @@ from sglang.multimodal_gen.configs.sample.diffusers_generic import (
     DiffusersGenericSamplingParams,
 )
 from sglang.multimodal_gen.configs.sample.ideogram import Ideogram4SamplingParams
+from sglang.multimodal_gen.configs.sample.pi05 import Pi05SamplingParams
 from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams
 
 __all__ = [
     "SamplingParams",
+    "VLASamplingParams",
     "DiffusersGenericSamplingParams",
     "Ideogram4SamplingParams",
+    "Pi05SamplingParams",
 ]

python/sglang/multimodal_gen/configs/sample/pi05.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams
+
+
+@dataclass
+class Pi05SamplingParams(VLASamplingParams):
+    """Sampling parameters for Pi0.5 flow-matching action inference."""
+
+    num_inference_steps: int = 10
+
+    action_horizon: int = 50
+    action_dim: int = 32
+    output_format: str = "list"
+    return_timing: bool = True
+    enable_prefix_cache: bool = True
+    enable_cuda_graph: bool = True
+
+    state: Any = field(default=None, metadata={"batch_sig_exclude": True})
+    images: dict[str, Any] | None = field(
+        default=None, metadata={"batch_sig_exclude": True}
+    )
+    image_masks: dict[str, bool] | None = field(
+        default=None, metadata={"batch_sig_exclude": True}
+    )
+    camera_order: list[str] | tuple[str, ...] | None = field(
+        default=None, metadata={"batch_sig_exclude": True}
+    )
+    noise: Any = field(default=None, metadata={"batch_sig_exclude": True})
+    observation: dict[str, Any] | None = field(
+        default=None, metadata={"batch_sig_exclude": True}
+    )
+
+    def build_request_extra(self) -> dict[str, Any]:
+        extra = super().build_request_extra()
+        observation = dict(self.observation or {})
+        if self.images is not None:
+            observation["images"] = self.images
+        if self.image_masks is not None:
+            observation["image_masks"] = self.image_masks
+        if self.state is not None:
+            observation["state"] = self.state
+        if self.camera_order is not None:
+            observation["camera_order"] = tuple(self.camera_order)
+        if self.prompt is not None:
+            observation["prompt"] = self.prompt
+        if self.noise is not None:
+            observation["noise"] = self.noise
+
+        extra["vla"] = {
+            "observation": observation,
+            "options": {
+                "output_format": self.output_format,
+                "return_timing": self.return_timing,
+                "enable_prefix_cache": self.enable_prefix_cache,
+                "enable_cuda_graph": self.enable_cuda_graph,
+            },
+        }
+        return extra
+
+    def _validate(self):
+        super()._validate()
+        if self.action_horizon <= 0:
+            raise ValueError("action_horizon must be positive")
+        if self.action_dim <= 0:
+            raise ValueError("action_dim must be positive")
+        if self.output_format not in ("list", "numpy"):
+            raise ValueError("output_format must be 'list' or 'numpy'")
+
+    def _set_output_file_name(self):
+        if self.output_file_name is None:
+            self.output_file_name = "pi05_action"
+        super()._set_output_file_name()

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -76,12 +76,15 @@ class DataType(Enum):
     IMAGE = auto()
     VIDEO = auto()
     MESH = auto()
+    ACTION = auto()
 
     def get_default_extension(self) -> str:
         if self == DataType.IMAGE:
             return "png"
         if self == DataType.VIDEO:
             return "mp4"
+        if self == DataType.ACTION:
+            return "json"
         return "glb"
 
 
@@ -246,10 +249,8 @@ class SamplingParams:
 
     def _set_output_file_ext(self):
         # add extension if needed
-        if not any(
-            self.output_file_name.endswith(ext)
-            for ext in [".mp4", ".jpg", ".png", ".webp", ".obj", ".glb"]
-        ):
+        output_extensions = (".mp4", ".jpg", ".png", ".webp", ".obj", ".glb", ".json")
+        if not any(self.output_file_name.endswith(ext) for ext in output_extensions):
             self.output_file_name = (
                 f"{self.output_file_name}.{self.data_type.get_default_extension()}"
             )
@@ -329,6 +330,8 @@ class SamplingParams:
 
     def _adjust_output_quality(self, output_quality: str, data_type: DataType) -> int:
         """Convert output_quality string to compression level."""
+        if data_type == DataType.ACTION:
+            return 0
         output_quality_mapper = {"maximum": 100, "high": 90, "medium": 55, "low": 35}
         if output_quality == "default":
             return 50 if data_type == DataType.VIDEO else 75
@@ -469,18 +472,22 @@ class SamplingParams:
         """
         check if the sampling params is compatible and valid with server_args
         """
-        if pipeline_config.task_type.requires_image_input():
+        task_type = pipeline_config.task_type
+        if task_type.is_action_gen():
+            return
+
+        if task_type.requires_image_input():
             # requires image input
             if self.image_path is None:
                 raise ValueError(
-                    f"Served model with task type '{pipeline_config.task_type.name}' requires an 'image_path' input, but none was provided"
+                    f"Served model with task type '{task_type.name}' requires an 'image_path' input, but none was provided"
                 )
 
-        if not pipeline_config.task_type.accepts_image_input():
+        if not task_type.accepts_image_input():
             # does not support image input
             if self.image_path is not None:
                 raise ValueError(
-                    f"input_reference is not supported for {pipeline_config.task_type.name} models."
+                    f"input_reference is not supported for {task_type.name} models."
                 )
 
     def _adjust(
@@ -494,7 +501,39 @@ class SamplingParams:
 
         # TODO: SamplingParams should not rely on ServerArgs
         pipeline_config = server_args.pipeline_config
+        task_type = pipeline_config.task_type
+        self.data_type = task_type.data_type()
 
+        self._adjust_output_path(server_args)
+        if task_type.is_action_gen():
+            self._adjust_action_fields(server_args)
+            return
+
+        if task_type.is_mesh_gen():
+            self._adjust_mesh_fields(server_args, pipeline_config)
+            return
+
+        if task_type.is_visual_gen():
+            self._adjust_visual_fields(server_args, pipeline_config)
+
+    def _adjust_output_path(self, server_args):
+        if self.output_path is None:
+            if server_args.output_path is not None:
+                self.output_path = server_args.output_path
+                logger.debug(
+                    f"Overriding output_path with server configuration: {self.output_path}"
+                )
+            else:
+                self.save_output = False
+
+    def _adjust_action_fields(self, server_args):
+        self.return_file_paths_only = False
+        self.num_frames = 1
+        self.adjust_frames = False
+        if self.save_output and not server_args.comfyui_mode:
+            self._set_output_file_name()
+
+    def _adjust_mesh_fields(self, server_args, pipeline_config):
         if self.guidance_scale is None:
             try:
                 from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
@@ -507,17 +546,16 @@ class SamplingParams:
                     self.guidance_scale = 1.0
             except ImportError:
                 self.guidance_scale = 1.0
+        self.return_frames = False
+        self.return_video = False
+        self.num_frames = 1
+        self.adjust_frames = False
+        if self.save_output and not server_args.comfyui_mode:
+            self._set_output_file_name()
 
-        self.data_type = server_args.pipeline_config.task_type.data_type()
-
-        if self.output_path is None:
-            if server_args.output_path is not None:
-                self.output_path = server_args.output_path
-                logger.debug(
-                    f"Overriding output_path with server configuration: {self.output_path}"
-                )
-            else:
-                self.save_output = False
+    def _adjust_visual_fields(self, server_args, pipeline_config):
+        if self.guidance_scale is None:
+            self.guidance_scale = 1.0
 
         # Process negative prompt
         if self.negative_prompt is not None and not self.negative_prompt.isspace():
@@ -576,7 +614,6 @@ class SamplingParams:
             if not server_args.pipeline_config.allow_set_num_frames():
                 logger.debug("Setting `num_frames` to 1 for image generation model")
                 self.num_frames = 1
-
         else:
             # mandatory frame adjusting logic, mod
             # NOTE: We must apply adjust_num_frames BEFORE the SP alignment logic below.

python/sglang/multimodal_gen/configs/sample/vla.py

@@ -0,0 +1,272 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import dataclasses
+import os
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    _sanitize_filename,
+)
+from sglang.multimodal_gen.utils import StoreBoolean, expand_path_fields
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+@dataclass
+class VLASamplingParams:
+    """Sampling parameters for VLA/action-generation policies."""
+
+    data_type: DataType = DataType.ACTION
+    request_id: str | None = field(default=None, metadata={"batch_sig_exclude": True})
+    prompt: str | list[str] | None = field(
+        default="", metadata={"batch_sig_exclude": True}
+    )
+    num_outputs_per_prompt: int = 1
+    seed: int | list[int] = field(default=42, metadata={"batch_sig_exclude": True})
+    generator_device: str | None = None
+    num_inference_steps: int = 10
+
+    output_path: str | None = field(default=None, metadata={"batch_sig_exclude": True})
+    output_file_name: str | None = field(
+        default=None, metadata={"batch_sig_exclude": True}
+    )
+    save_output: bool = False
+    return_file_paths_only: bool = False
+
+    profile: bool = field(default=False, metadata={"batch_sig_exclude": True})
+    num_profiled_timesteps: int = field(default=5, metadata={"batch_sig_exclude": True})
+    profile_all_stages: bool = field(
+        default=False, metadata={"batch_sig_exclude": True}
+    )
+    debug: bool = field(default=False, metadata={"batch_sig_exclude": True})
+    perf_dump_path: str | None = field(
+        default=None, metadata={"batch_sig_exclude": True}
+    )
+    suppress_logs: bool = field(default=False, metadata={"batch_sig_exclude": True})
+
+    enable_sequence_shard: bool | None = None
+    max_sequence_length: int | None = None
+    no_override_protected_fields: bool = field(
+        default=False, metadata={"batch_sig_exclude": True}
+    )
+
+    def __post_init__(self) -> None:
+        self.data_type = DataType.ACTION
+        self._validate()
+
+        env_steps = os.environ.get("SGLANG_TEST_NUM_INFERENCE_STEPS")
+        if env_steps is not None and self.num_inference_steps is not None:
+            self.num_inference_steps = int(env_steps)
+
+    def build_request_extra(self) -> dict[str, Any]:
+        extra = {}
+        diffusers_kwargs = getattr(self, "diffusers_kwargs", None)
+        if diffusers_kwargs:
+            extra["diffusers_kwargs"] = diffusers_kwargs
+        explicit_fields = getattr(self, "_explicit_fields", None)
+        if explicit_fields is not None:
+            extra["explicit_fields"] = sorted(explicit_fields)
+        return extra
+
+    def apply_request_extra(self, req: Any) -> None:
+        req.extra.update(self.build_request_extra())
+
+    def _validate(self):
+        if (
+            not isinstance(self.num_outputs_per_prompt, int)
+            or self.num_outputs_per_prompt <= 0
+        ):
+            raise ValueError(
+                "num_outputs_per_prompt must be a positive int, "
+                f"got {self.num_outputs_per_prompt!r}"
+            )
+
+        if isinstance(self.seed, list):
+            if not self.seed:
+                raise ValueError("seed list must not be empty")
+            for seed in self.seed:
+                if isinstance(seed, bool) or not isinstance(seed, int) or seed < 0:
+                    raise ValueError(
+                        f"seed list must contain non-negative ints, got {self.seed!r}"
+                    )
+        elif (
+            isinstance(self.seed, bool)
+            or not isinstance(self.seed, int)
+            or self.seed < 0
+        ):
+            raise ValueError(
+                f"seed must be a non-negative int or list of ints, got {self.seed!r}"
+            )
+
+        if (
+            not isinstance(self.num_inference_steps, int)
+            or self.num_inference_steps <= 0
+        ):
+            raise ValueError(
+                "num_inference_steps must be a positive int, "
+                f"got {self.num_inference_steps!r}"
+            )
+
+        if self.generator_device not in (None, "cuda", "musa", "cpu"):
+            raise ValueError(
+                "generator_device must be one of None, 'cuda', 'musa', or 'cpu', "
+                f"got {self.generator_device!r}"
+            )
+
+    def _validate_with_pipeline_config(self, pipeline_config):
+        if not pipeline_config.task_type.is_action_gen():
+            raise ValueError(
+                f"VLASamplingParams requires an ACTION pipeline, got {pipeline_config.task_type.name}"
+            )
+
+    def _adjust(self, server_args: "ServerArgs"):
+        expand_path_fields(self)
+        self.data_type = DataType.ACTION
+        self.return_file_paths_only = False
+        if self.output_path is None and server_args.output_path is not None:
+            self.output_path = server_args.output_path
+        if self.output_path is None:
+            self.save_output = False
+        if self.save_output and not server_args.comfyui_mode:
+            self._set_output_file_name()
+
+    def _set_output_file_ext(self):
+        if self.output_file_name and not self.output_file_name.endswith(".json"):
+            self.output_file_name = f"{self.output_file_name}.json"
+
+    def _set_output_file_name(self):
+        if self.output_file_name is None:
+            self.output_file_name = "vla_action"
+        self.output_file_name = _sanitize_filename(self.output_file_name)
+        self._set_output_file_ext()
+
+    def output_file_path(self):
+        if self.output_path is None or self.output_file_name is None:
+            return None
+        return os.path.join(self.output_path, self.output_file_name)
+
+    def _merge_with_user_params(
+        self,
+        user_params: "VLASamplingParams",
+        explicit_fields: set[str] | None = None,
+    ):
+        if user_params is None:
+            return
+
+        predefined_fields = set(type(self).__annotations__.keys())
+        allow_override_protected = not user_params.no_override_protected_fields
+        for field_info in dataclasses.fields(user_params):
+            field_name = field_info.name
+            user_value = getattr(user_params, field_name)
+            if field_info.default is not dataclasses.MISSING:
+                default_class_value = field_info.default
+            elif field_info.default_factory is not dataclasses.MISSING:
+                default_class_value = field_info.default_factory()
+            else:
+                default_class_value = dataclasses.MISSING
+
+            if explicit_fields is not None:
+                is_user_modified = field_name in explicit_fields
+            else:
+                is_user_modified = user_value != default_class_value
+            is_protected_field = field_name in predefined_fields
+            if is_user_modified and (
+                allow_override_protected or not is_protected_field
+            ):
+                setattr(self, field_name, user_value)
+
+        if explicit_fields is not None:
+            self._explicit_fields = set(explicit_fields)
+        self.__post_init__()
+
+    @staticmethod
+    def add_cli_args(parser: Any) -> Any:
+        def add_argument(*name_or_flags, **kwargs):
+            kwargs.setdefault("default", argparse.SUPPRESS)
+            return parser.add_argument(*name_or_flags, **kwargs)
+
+        add_argument(
+            "--prompt",
+            type=str,
+            nargs="+",
+            help="Language instruction(s) for the VLA policy.",
+        )
+        add_argument(
+            "--num-inference-steps",
+            type=int,
+            help="Number of action denoising steps.",
+        )
+        add_argument(
+            "--num-outputs-per-prompt",
+            type=int,
+            help="Number of candidate actions to generate per observation.",
+        )
+        add_argument(
+            "--seed",
+            type=int,
+            nargs="+",
+            help="Random seed for action noise generation.",
+        )
+        add_argument(
+            "--generator-device",
+            type=str,
+            choices=["cuda", "musa", "cpu"],
+            help="Device for random generator. Default: use the model-specific setting.",
+        )
+        add_argument(
+            "--profile",
+            action="store_true",
+            help="Enable torch profiler for action denoising.",
+        )
+        add_argument(
+            "--num-profiled-timesteps",
+            type=int,
+            help="Number of denoising timesteps to profile after warmup.",
+        )
+        add_argument(
+            "--profile-all-stages",
+            action="store_true",
+            dest="profile_all_stages",
+            help="Used with --profile, profile all pipeline stages.",
+        )
+        add_argument("--debug", action="store_true")
+        add_argument(
+            "--enable-sequence-shard",
+            action=StoreBoolean,
+            help="Enable sequence dimension shard with sequence parallelism.",
+        )
+        add_argument(
+            "--max-sequence-length",
+            type=int,
+            help="Maximum prefix sequence length.",
+        )
+        add_argument(
+            "--no-override-protected-fields",
+            action="store_true",
+            help="If set, disallow user params to override subclass-defined fields.",
+        )
+        return parser
+
+    @classmethod
+    def get_cli_args(cls, args: argparse.Namespace):
+        sampling_params_fields = {attr.name for attr in dataclasses.fields(cls)}
+        args_attrs = set(vars(args).keys())
+        attrs = sampling_params_fields & args_attrs
+        cli_args = {
+            attr: getattr(args, attr)
+            for attr in attrs
+            if hasattr(args, attr) and getattr(args, attr) is not None
+        }
+        if isinstance(cli_args.get("seed"), list) and len(cli_args["seed"]) == 1:
+            cli_args["seed"] = cli_args["seed"][0]
+        return cli_args
+
+    def output_size_str(self) -> str:
+        return "action"
+
+    def seconds(self) -> float:
+        return 0.0

python/sglang/multimodal_gen/registry.py

@@ -76,6 +76,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.mova import (
     MOVA360PConfig,
     MOVA720PConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     QwenImageEditPipelineConfig,
     QwenImageEditPlus_2511_PipelineConfig,
@@ -137,6 +138,7 @@ from sglang.multimodal_gen.configs.sample.mova import (
     MOVA_360P_SamplingParams,
     MOVA_720P_SamplingParams,
 )
+from sglang.multimodal_gen.configs.sample.pi05 import Pi05SamplingParams
 from sglang.multimodal_gen.configs.sample.qwenimage import (
     QwenImage2512SamplingParams,
     QwenImageEditPlusSamplingParams,
@@ -629,6 +631,20 @@ def get_model_info(
 
 # Registration of model configs
 def _register_configs():
+    # Pi0.5 / OpenPI / LeRobot action policies.
+    register_configs(
+        sampling_param_cls=Pi05SamplingParams,
+        pipeline_config_cls=Pi05PipelineConfig,
+        hf_model_paths=[
+            "lerobot/pi05_base",
+            "lerobot/pi05_libero_base",
+        ],
+        model_detectors=[
+            lambda hf_id: "pi05" in hf_id.lower(),
+            lambda hf_id: "pi0.5" in hf_id.lower(),
+        ],
+    )
+
     # LTX-2
     register_configs(
         sampling_param_cls=LTX2SamplingParams,

python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py

@@ -602,10 +602,11 @@ class GroupCoordinator:
         group = self.device_group
         metadata_group = self.cpu_group
         assert src < self.world_size, f"Invalid src rank ({src})"
-        src = self.ranks[src]
+        src_rank_in_group = src
+        src_global_rank = self.ranks[src_rank_in_group]
 
         rank = self.rank
-        if rank == src:
+        if rank == src_global_rank:
             metadata_list: List[Tuple[Any, Any]] = []
             assert isinstance(
                 tensor_dict, dict
@@ -614,7 +615,7 @@ class GroupCoordinator:
             # `metadata_list` lives in CPU memory.
             # `broadcast_object_list` has serialization & deserialization,
             # all happening on CPU. Therefore, we can use the CPU group.
-            self.broadcast_object(metadata_list, src=src)
+            self.broadcast_object(metadata_list, src=src_rank_in_group)
             async_handles = []
             for tensor in tensor_list:
                 if tensor.numel() == 0:
@@ -623,19 +624,22 @@ class GroupCoordinator:
                 if tensor.is_cpu:
                     # use metadata_group for CPU tensors
                     handle = torch.distributed.broadcast(
-                        tensor, src=src, group=metadata_group, async_op=True
+                        tensor,
+                        src=src_global_rank,
+                        group=metadata_group,
+                        async_op=True,
                     )
                 else:
                     # use group for GPU tensors
                     handle = torch.distributed.broadcast(
-                        tensor, src=src, group=group, async_op=True
+                        tensor, src=src_global_rank, group=group, async_op=True
                     )
                 async_handles.append(handle)
             for async_handle in async_handles:
                 async_handle.wait()
 
         else:
-            metadata_list = self.broadcast_object(None, src=src)
+            metadata_list = self.broadcast_object(None, src=src_rank_in_group)
             tensor_dict = {}
             async_handles = []
             for key, value in metadata_list:
@@ -650,12 +654,15 @@ class GroupCoordinator:
                     if tensor.is_cpu:
                         # use metadata_group for CPU tensors
                         handle = torch.distributed.broadcast(
-                            tensor, src=src, group=metadata_group, async_op=True
+                            tensor,
+                            src=src_global_rank,
+                            group=metadata_group,
+                            async_op=True,
                         )
                     else:
                         # use group for GPU tensors
                         handle = torch.distributed.broadcast(
-                            tensor, src=src, group=group, async_op=True
+                            tensor, src=src_global_rank, group=group, async_op=True
                         )
                     async_handles.append(handle)
                     _update_nested_dict(tensor_dict, key, tensor)

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -376,6 +376,34 @@ class DiffGenerator:
             return None
         return results[0] if len(results) == 1 else results
 
+    def generate_action(
+        self,
+        sampling_params_kwargs: dict | None = None,
+        external_trace_header: dict[str, str] | None = None,
+    ) -> dict[str, Any]:
+        sampling_params_kwargs = sampling_params_kwargs or {}
+        sampling_params = SamplingParams.from_user_sampling_params_args(
+            self.server_args.model_path,
+            server_args=self.server_args,
+            **sampling_params_kwargs,
+        )
+        if sampling_params.data_type != DataType.ACTION:
+            raise ValueError(
+                f"generate_action requires an ACTION pipeline, got {sampling_params.data_type}"
+            )
+
+        req = prepare_request(
+            server_args=self.server_args,
+            sampling_params=sampling_params,
+            external_trace_header=external_trace_header,
+        )
+        output_batch = self._send_to_scheduler_and_wait_for_response(req)
+        if output_batch.error:
+            raise RuntimeError(output_batch.error)
+        if output_batch.output is None:
+            raise RuntimeError("action policy returned no output")
+        return output_batch.output[0]
+
     def _resolve_prompts(
         self,
         prompt: str | list[str] | None,
@@ -430,9 +458,13 @@ class DiffGenerator:
             and output_index < len(output_batch.metrics_list)
         ):
             metrics = output_batch.metrics_list[output_index]
+        if req.data_type == DataType.ACTION:
+            size = ("action",)
+        else:
+            size = (req.height, req.width, req.num_frames)
         return dict(
             prompt=req.prompt,
-            size=(req.height, req.width, req.num_frames),
+            size=size,
             generation_time=generation_time,
             peak_memory_mb=output_batch.peak_memory_mb,
             metrics=metrics.to_dict() if metrics else {},

python/sglang/multimodal_gen/runtime/entrypoints/http_server.py

@@ -14,10 +14,7 @@ from fastapi import APIRouter, FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 
 from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
-from sglang.multimodal_gen.runtime.entrypoints.openai import (
-    image_api,
-    video_api,
-)
+from sglang.multimodal_gen.runtime.entrypoints.openai import image_api, video_api
 from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
     VertexGenerateReqInput,
 )
@@ -33,6 +30,8 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
     prepare_request,
     save_outputs,
 )
+from sglang.multimodal_gen.runtime.entrypoints.vla import api as vla_api
+from sglang.multimodal_gen.runtime.entrypoints.vla import openpi
 from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
 from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
 from sglang.multimodal_gen.runtime.server_warmup import (
@@ -403,6 +402,9 @@ def create_app(server_args: ServerArgs):
     app.include_router(image_api.router)
     app.include_router(video_api.router)
     app.include_router(realtime_video_api.router)
+    if server_args.pipeline_config.task_type.is_action_gen():
+        app.include_router(vla_api.router)
+        app.include_router(openpi.router)
     app.include_router(mesh_api.router)
     app.include_router(weights_api.router)
     app.include_router(rollout_api.router)

python/sglang/multimodal_gen/runtime/entrypoints/utils.py

@@ -8,6 +8,7 @@ This module provides a consolidated interface for generating videos using
 diffusion models.
 """
 
+import json
 import os
 import shutil
 import subprocess
@@ -445,11 +446,13 @@ def prepare_request(
     if not isinstance(req.prompt, str):
         raise TypeError(f"`prompt` must be a string, but got {type(req.prompt)}")
 
-    if (req.width is not None and req.width <= 0) or (
-        req.height is not None and req.height <= 0
+    req_width = getattr(req, "width", None)
+    req_height = getattr(req, "height", None)
+    if (req_width is not None and req_width <= 0) or (
+        req_height is not None and req_height <= 0
     ):
         raise ValueError(
-            f"Height and width must be positive, got height={req.height}, width={req.width}"
+            f"Height and width must be positive, got height={req_height}, width={req_width}"
         )
 
     if server_args.enable_trace:
@@ -661,6 +664,21 @@ def save_outputs(
     output_paths: list[str] = []
     for idx, sample in enumerate(outputs):
         save_file_path = build_output_path(idx)
+        if data_type == DataType.ACTION:
+            if samples_out is not None:
+                samples_out.append(sample)
+            if audios_out is not None:
+                audios_out.append(None)
+            if frames_out is not None:
+                frames_out.append([])
+            if save_output and save_file_path:
+                os.makedirs(os.path.dirname(save_file_path) or ".", exist_ok=True)
+                with open(save_file_path, "w", encoding="utf-8") as f:
+                    json.dump(sample, f, ensure_ascii=False)
+                logger.info(f"Output saved to {CYAN}{save_file_path}{RESET}")
+            output_paths.append(save_file_path)
+            continue
+
         if data_type == DataType.VIDEO:
             sample = attach_audio_to_video_sample(sample, audio, idx)
 
@@ -711,6 +729,9 @@ def post_process_sample(
     upscaling_scale: int = 4,
 ) -> list[Any]:
     """materialize frames and save outputs (optional)"""
+    if data_type == DataType.ACTION:
+        return []
+
     materialized = materialize_output_sample(
         sample,
         data_type,

python/sglang/multimodal_gen/runtime/entrypoints/vla/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

python/sglang/multimodal_gen/runtime/entrypoints/vla/api.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from fastapi import APIRouter, HTTPException, Request, Response, WebSocket
+
+from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (
+    action_generation_response,
+    action_metadata,
+    action_raw_response,
+    infer_action,
+    pack_msgpack,
+    unpack_msgpack,
+)
+from sglang.multimodal_gen.runtime.entrypoints.vla.ws_utils import (
+    run_action_msgpack_ws,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.srt.utils.json_response import orjson_response
+
+router = APIRouter(prefix="/v1/actions", tags=["actions"])
+
+
+def _wants_msgpack(request: Request) -> bool:
+    content_type = request.headers.get("content-type", "").lower()
+    accept = request.headers.get("accept", "").lower()
+    return "msgpack" in content_type or "msgpack" in accept
+
+
+def _response_format(payload: dict) -> str:
+    runtime = payload.get("runtime") or {}
+    response_format = str(runtime.get("response_format", "envelope")).lower()
+    if response_format not in ("envelope", "raw"):
+        raise ValueError("runtime.response_format must be 'envelope' or 'raw'")
+    return response_format
+
+
+def _prefer_numpy_output(payload: dict) -> None:
+    runtime = payload.setdefault("runtime", {})
+    runtime.setdefault("output_format", "numpy")
+
+
+@router.post("/generations")
+async def create_action_generation(request: Request):
+    server_args: ServerArgs = request.app.state.server_args
+    try:
+        if "msgpack" in request.headers.get("content-type", "").lower():
+            payload = unpack_msgpack(await request.body())
+        else:
+            payload = await request.json()
+        wants_msgpack = _wants_msgpack(request)
+        if wants_msgpack:
+            _prefer_numpy_output(payload)
+        output = await infer_action(payload, server_args)
+        if _response_format(payload) == "raw":
+            response = action_raw_response(output, preserve_numpy=wants_msgpack)
+        else:
+            response = action_generation_response(
+                output,
+                server_args,
+                preserve_numpy=wants_msgpack,
+            )
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+    if wants_msgpack:
+        return Response(
+            content=pack_msgpack(response), media_type="application/msgpack"
+        )
+    return orjson_response(response)
+
+
+@router.get("/metadata")
+async def action_metadata_endpoint(request: Request):
+    return orjson_response(action_metadata(request.app.state.server_args))
+
+
+@router.websocket("/realtime")
+async def action_realtime_ws(websocket: WebSocket):
+    server_args: ServerArgs = websocket.app.state.server_args
+    await run_action_msgpack_ws(
+        websocket,
+        server_args,
+        prepare_payload=_prefer_numpy_output,
+        build_response=lambda output: action_generation_response(
+            output,
+            server_args,
+            preserve_numpy=True,
+        ),
+    )

python/sglang/multimodal_gen/runtime/entrypoints/vla/openpi.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import APIRouter, WebSocket
+
+from sglang.multimodal_gen.runtime.entrypoints.vla.ws_utils import (
+    run_action_msgpack_ws,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+router = APIRouter()
+
+
+def _prefer_numpy_output(observation: dict[str, Any]) -> None:
+    observation.setdefault("output_format", "numpy")
+
+
+@router.websocket("/openpi/policy")
+async def openpi_policy_ws(websocket: WebSocket):
+    server_args: ServerArgs = websocket.app.state.server_args
+    await run_action_msgpack_ws(
+        websocket,
+        server_args,
+        prepare_payload=_prefer_numpy_output,
+        build_response=lambda output: output,
+    )

python/sglang/multimodal_gen/runtime/entrypoints/vla/protocol.py

@@ -0,0 +1,443 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import base64
+import dataclasses
+import io
+import time
+import uuid
+from functools import lru_cache
+from typing import Any
+
+import numpy as np
+from PIL import Image
+
+from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+def pack_numpy_payload(obj):
+    if isinstance(obj, (np.ndarray, np.generic)) and obj.dtype.kind in ("V", "O", "c"):
+        raise ValueError(f"Unsupported dtype: {obj.dtype}")
+    if isinstance(obj, np.ndarray):
+        return {
+            b"__ndarray__": True,
+            b"data": obj.tobytes(),
+            b"dtype": obj.dtype.str,
+            b"shape": obj.shape,
+        }
+    if isinstance(obj, np.generic):
+        return {
+            b"__npgeneric__": True,
+            b"data": obj.item(),
+            b"dtype": obj.dtype.str,
+        }
+    return obj
+
+
+def unpack_numpy_payload(obj):
+    ndarray_marker = obj.get("__ndarray__") or obj.get(b"__ndarray__")
+    npgeneric_marker = obj.get("__npgeneric__") or obj.get(b"__npgeneric__")
+    data = obj.get("data", obj.get(b"data"))
+    dtype = obj.get("dtype", obj.get(b"dtype"))
+    shape = obj.get("shape", obj.get(b"shape"))
+    if ndarray_marker:
+        return np.ndarray(
+            buffer=data,
+            dtype=np.dtype(dtype),
+            shape=shape,
+        )
+    if npgeneric_marker:
+        return np.dtype(dtype).type(data)
+    return obj
+
+
+def pack_msgpack(payload: Any) -> bytes:
+    import msgpack
+
+    return msgpack.packb(payload, default=pack_numpy_payload, use_bin_type=True)
+
+
+def unpack_msgpack(payload: bytes) -> Any:
+    import msgpack
+
+    return msgpack.unpackb(payload, object_hook=unpack_numpy_payload, raw=False)
+
+
+def _decode_b64_image(payload: dict[str, Any]) -> Image.Image:
+    data = payload.get("b64_json") or payload.get("base64")
+    if not data:
+        raise ValueError("image payload requires b64_json")
+    if isinstance(data, str) and "," in data and data.startswith("data:"):
+        data = data.split(",", 1)[1]
+    return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGB")
+
+
+def _decode_tensor_payload(payload: dict[str, Any]) -> Any:
+    values = payload.get("values")
+    if values is None:
+        values = payload.get("data")
+    if values is None:
+        return payload
+    dtype = payload.get("dtype")
+    array = np.asarray(values, dtype=np.dtype(dtype) if dtype else None)
+    shape = payload.get("shape")
+    if shape is not None:
+        array = array.reshape(tuple(shape))
+    return array
+
+
+def _normalize_image_value(value: Any) -> Any:
+    if not isinstance(value, dict):
+        return value
+    if "b64_json" in value or "base64" in value:
+        return _decode_b64_image(value)
+    if "values" in value or "data" in value:
+        return _decode_tensor_payload(value)
+    return value
+
+
+def _normalize_observation(observation: dict[str, Any]) -> dict[str, Any]:
+    normalized = dict(observation)
+    images = normalized.get("images")
+    if isinstance(images, dict):
+        normalized["images"] = {
+            name: _normalize_image_value(value) for name, value in images.items()
+        }
+    state = normalized.get("state")
+    if isinstance(state, dict):
+        normalized["state"] = _decode_tensor_payload(state)
+    observation_state = normalized.get("observation.state")
+    if isinstance(observation_state, dict):
+        normalized["observation.state"] = _decode_tensor_payload(observation_state)
+    noise = normalized.get("noise")
+    if isinstance(noise, dict):
+        normalized["noise"] = _decode_tensor_payload(noise)
+    observation_noise = normalized.get("observation.noise")
+    if isinstance(observation_noise, dict):
+        normalized["observation.noise"] = _decode_tensor_payload(observation_noise)
+    return normalized
+
+
+def images_from_observation(
+    observation: dict[str, Any],
+    pipeline_config: Any,
+) -> dict[str, Any]:
+    if isinstance(observation.get("images"), dict):
+        images = dict(observation["images"])
+    else:
+        images = {}
+    for key in pipeline_config.image_keys:
+        if key in observation:
+            images[key] = observation[key]
+        full_key = f"observation.images.{key}"
+        if full_key in observation:
+            images[key] = observation[full_key]
+    return {name: _normalize_image_value(value) for name, value in images.items()}
+
+
+def action_metadata(server_args: ServerArgs) -> dict[str, Any]:
+    pipeline_config = server_args.pipeline_config
+    policy_family = getattr(
+        pipeline_config,
+        "policy_family",
+        type(pipeline_config).__name__.removesuffix("PipelineConfig").lower(),
+    )
+    return {
+        "object": "action.metadata",
+        "model": server_args.model_id or server_args.model_path,
+        "model_path": server_args.model_path,
+        "policy_family": policy_family,
+        "input": {
+            "image_keys": list(pipeline_config.image_keys),
+            "image_size": list(pipeline_config.image_size),
+            "state_dim": pipeline_config.state_dim,
+        },
+        "output": {
+            "action_type": "continuous",
+            "action_horizon": pipeline_config.action_horizon,
+            "action_dim": pipeline_config.output_action_dim,
+            "padded_action_dim": pipeline_config.action_dim,
+            "dtype": "float32",
+        },
+        "runtime": {
+            "materialize_dtype": pipeline_config.materialize_dtype,
+            "enable_autocast": pipeline_config.enable_autocast,
+            "parallelism": {
+                "num_gpus": server_args.num_gpus,
+                "tp_size": server_args.tp_size,
+                "sp_degree": server_args.sp_degree,
+                "ulysses_degree": server_args.ulysses_degree,
+                "ring_degree": server_args.ring_degree,
+                "prefix_strategy": pipeline_config.prefix_parallel_strategy,
+                "action_strategy": pipeline_config.action_parallel_strategy,
+                "layout_version": pipeline_config.parallel_layout_version,
+            },
+        },
+        "defaults": {
+            "num_inference_steps": pipeline_config.default_num_inference_steps,
+            "prefix_cache": (
+                "auto" if pipeline_config.enable_global_prefix_cache else False
+            ),
+            "cuda_graph": "auto" if pipeline_config.enable_action_cuda_graph else False,
+        },
+        "capabilities": {
+            "exact_prefix_cache": True,
+            "cuda_graph": pipeline_config.enable_action_cuda_graph,
+            "realtime_websocket": True,
+            "openpi_websocket": True,
+            "batch_inputs": False,
+            "multiple_candidates": False,
+        },
+    }
+
+
+def _runtime_bool(value: Any, default: bool) -> bool:
+    if value is None:
+        return default
+    if isinstance(value, str):
+        value = value.lower()
+        if value == "auto":
+            return default
+        if value in ("true", "1", "yes"):
+            return True
+        if value in ("false", "0", "no"):
+            return False
+    return bool(value)
+
+
+def _action_request_to_observation(payload: dict[str, Any]) -> dict[str, Any]:
+    if "input" not in payload:
+        return _normalize_observation(payload)
+
+    input_payload = payload.get("input") or {}
+    observation = dict(input_payload.get("observation") or {})
+    if "task" in input_payload:
+        observation["prompt"] = input_payload["task"]
+    elif "prompt" in input_payload:
+        observation["prompt"] = input_payload["prompt"]
+    if "images" in input_payload:
+        observation["images"] = input_payload["images"]
+    if "state" in input_payload:
+        observation["state"] = input_payload["state"]
+    if "noise" in input_payload:
+        observation["noise"] = input_payload["noise"]
+    return _normalize_observation(observation)
+
+
+@lru_cache(maxsize=32)
+def _resolve_action_sampling_params_cls_cached(
+    model_path: str,
+    backend: str | None,
+    model_id: str | None,
+    pipeline_class_name: str | None,
+) -> type[VLASamplingParams]:
+    if pipeline_class_name:
+        from sglang.multimodal_gen.registry import get_pipeline_config_classes
+
+        config_classes = get_pipeline_config_classes(pipeline_class_name)
+        if config_classes is not None:
+            _, sampling_params_cls = config_classes
+            if issubclass(sampling_params_cls, VLASamplingParams):
+                return sampling_params_cls
+
+    from sglang.multimodal_gen.registry import get_model_info
+
+    model_info = get_model_info(
+        model_path,
+        backend=backend,
+        model_id=model_id,
+    )
+    sampling_params_cls = model_info.sampling_param_cls
+    if not issubclass(sampling_params_cls, VLASamplingParams):
+        raise ValueError(
+            f"Action endpoint requires VLASamplingParams, got {sampling_params_cls.__name__}"
+        )
+    return sampling_params_cls
+
+
+def _resolve_action_sampling_params_cls(
+    server_args: ServerArgs,
+) -> type[VLASamplingParams]:
+    return _resolve_action_sampling_params_cls_cached(
+        server_args.model_path,
+        getattr(server_args, "backend", None),
+        getattr(server_args, "model_id", None),
+        getattr(server_args, "pipeline_class_name", None),
+    )
+
+
+@lru_cache(maxsize=32)
+def _sampling_params_field_names(
+    sampling_params_cls: type[VLASamplingParams],
+) -> frozenset[str]:
+    return frozenset(field.name for field in dataclasses.fields(sampling_params_cls))
+
+
+def build_action_sampling_params(
+    payload: dict[str, Any],
+    server_args: ServerArgs,
+) -> VLASamplingParams:
+    pipeline_config = server_args.pipeline_config
+    observation = _action_request_to_observation(payload)
+    parameters = dict(payload.get("parameters") or {})
+    runtime = dict(payload.get("runtime") or {})
+    if "return_timing" in payload and "return_timing" not in runtime:
+        runtime["return_timing"] = payload["return_timing"]
+    images = images_from_observation(observation, pipeline_config)
+    state = observation.get("state")
+    if state is None:
+        state = observation.get("observation.state")
+    noise = observation.get("noise")
+    if noise is None:
+        noise = observation.get("observation.noise")
+    prompt = observation.get("prompt") or observation.get("task") or ""
+    prefix_cache = runtime.get("prefix_cache")
+    if prefix_cache is None:
+        prefix_cache = observation.get("enable_prefix_cache")
+    if prefix_cache is None:
+        prefix_cache = observation.get("enable_pi_prefix_cache")
+    cuda_graph = runtime.get("cuda_graph")
+    if cuda_graph is None:
+        cuda_graph = observation.get("enable_cuda_graph")
+    if cuda_graph is None:
+        cuda_graph = observation.get("enable_pi_cuda_graph")
+    output_format = str(
+        runtime.get(
+            "output_format",
+            parameters.get(
+                "output_format",
+                observation.get("output_format", "list"),
+            ),
+        )
+    ).lower()
+    if output_format not in ("list", "numpy"):
+        raise ValueError("output_format must be 'list' or 'numpy'")
+
+    sampling_params_cls = _resolve_action_sampling_params_cls(server_args)
+    sampling_kwargs = {
+        "request_id": payload.get("request_id") or payload.get("id"),
+        "prompt": prompt,
+        "images": images,
+        "image_masks": observation.get("image_masks"),
+        "camera_order": observation.get("camera_order"),
+        "state": state,
+        "noise": noise,
+        "observation": observation,
+        "action_horizon": int(
+            parameters.get(
+                "action_horizon",
+                observation.get("action_horizon", pipeline_config.action_horizon),
+            )
+        ),
+        "action_dim": int(
+            parameters.get(
+                "action_dim",
+                observation.get("action_dim", pipeline_config.action_dim),
+            )
+        ),
+        "num_inference_steps": int(
+            parameters.get(
+                "num_inference_steps",
+                observation.get(
+                    "num_inference_steps",
+                    pipeline_config.default_num_inference_steps,
+                ),
+            )
+        ),
+        "output_format": output_format,
+        "return_timing": _runtime_bool(runtime.get("return_timing"), True),
+        "enable_prefix_cache": _runtime_bool(prefix_cache, True),
+        "enable_cuda_graph": _runtime_bool(cuda_graph, True),
+    }
+    supported_fields = _sampling_params_field_names(sampling_params_cls)
+    sp = sampling_params_cls(
+        **{
+            name: value
+            for name, value in sampling_kwargs.items()
+            if name in supported_fields
+        }
+    )
+    sp._adjust(server_args)
+    return sp
+
+
+async def infer_action(
+    payload: dict[str, Any],
+    server_args: ServerArgs,
+) -> dict[str, Any]:
+    sp = build_action_sampling_params(payload, server_args)
+    req = prepare_request(server_args, sp)
+    response = await async_scheduler_client.forward(req)
+    if getattr(response, "error", None):
+        raise RuntimeError(response.error)
+    if response.output is None:
+        raise RuntimeError("action policy returned no output")
+    return response.output[0]
+
+
+def action_generation_response(
+    output: dict[str, Any],
+    server_args: ServerArgs,
+    *,
+    preserve_numpy: bool = False,
+) -> dict[str, Any]:
+    actions = output["actions"]
+    if isinstance(actions, np.ndarray):
+        action_shape = list(actions.shape)
+        action_values = actions if preserve_numpy else actions.tolist()
+    else:
+        horizon = len(actions) if isinstance(actions, list) else 0
+        action_dim = len(actions[0]) if horizon and isinstance(actions[0], list) else 0
+        action_shape = [horizon, action_dim]
+        action_values = actions
+    response = {
+        "id": output.get("request_id") or f"act_{uuid.uuid4().hex}",
+        "object": "action.generation",
+        "created": int(time.time()),
+        "model": server_args.model_id or server_args.model_path,
+        "data": [
+            {
+                "index": 0,
+                "input_index": 0,
+                "candidate_index": 0,
+                "action": {
+                    "type": "continuous",
+                    "dtype": "float32",
+                    "shape": action_shape,
+                    "values": action_values,
+                },
+            }
+        ],
+        "usage": {
+            "action_horizon": action_shape[0] if action_shape else 0,
+            "action_dim": action_shape[1] if len(action_shape) > 1 else 0,
+            "denoise_steps": output.get("parameters", {}).get(
+                "num_inference_steps",
+                server_args.pipeline_config.default_num_inference_steps,
+            ),
+            "prefix_cache_hit": bool(output.get("cache", {}).get("hit", False)),
+        },
+    }
+    if "timings" in output:
+        response["timings"] = output["timings"]
+    if "cache" in output:
+        response["cache"] = output["cache"]
+    if "parallel" in output:
+        response["parallel"] = output["parallel"]
+    return response
+
+
+def action_raw_response(
+    output: dict[str, Any],
+    *,
+    preserve_numpy: bool = False,
+) -> dict[str, Any]:
+    response = dict(output)
+    actions = response.get("actions")
+    if isinstance(actions, np.ndarray) and not preserve_numpy:
+        response["actions"] = actions.tolist()
+    return response

python/sglang/multimodal_gen/runtime/entrypoints/vla/ws_utils.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import time
+import traceback
+from collections.abc import Callable
+from typing import Any
+
+from fastapi import WebSocket, WebSocketDisconnect
+
+from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (
+    action_metadata,
+    infer_action,
+    pack_msgpack,
+    unpack_msgpack,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+async def run_action_msgpack_ws(
+    websocket: WebSocket,
+    server_args: ServerArgs,
+    *,
+    prepare_payload: Callable[[dict[str, Any]], None],
+    build_response: Callable[[dict[str, Any]], dict[str, Any]],
+) -> None:
+    await websocket.accept()
+    await websocket.send_bytes(pack_msgpack(action_metadata(server_args)))
+
+    prev_total_time = None
+    while True:
+        try:
+            start_time = time.monotonic()
+            payload = unpack_msgpack(await websocket.receive_bytes())
+            prepare_payload(payload)
+            infer_start = time.monotonic()
+            output = await infer_action(payload, server_args)
+            response = build_response(output)
+            response.setdefault("server_timing", {})["infer_ms"] = (
+                time.monotonic() - infer_start
+            ) * 1000
+            if prev_total_time is not None:
+                response["server_timing"]["prev_total_ms"] = prev_total_time * 1000
+            await websocket.send_bytes(pack_msgpack(response))
+            prev_total_time = time.monotonic() - start_time
+        except WebSocketDisconnect:
+            break
+        except Exception:
+            try:
+                await websocket.send_bytes(
+                    pack_msgpack({"error": traceback.format_exc()})
+                )
+            except Exception:
+                pass
+            await websocket.close(code=1011, reason="Internal server error")
+            raise

python/sglang/multimodal_gen/runtime/launch_server.py

@@ -279,6 +279,12 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True):
     logger.debug("All workers are ready")
 
     if launch_http_server:
+        if server_args.pipeline_config.task_type.is_action_gen():
+            logger.info(
+                "VLA pipeline ready: model=%s; per-request details are "
+                "debug-only (use --log-level debug).",
+                server_args.model_id or server_args.model_path,
+            )
         logger.info("Starting FastAPI server.")
         if server_args.webui:
             logger.info("Launch FastAPI server in another process because of webui.")

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -420,6 +420,7 @@ class LocalAttention(nn.Module):
         softmax_scale: float | None = None,
         causal: bool = False,
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        compute_dtype: torch.dtype | None = None,
         **extra_impl_args,
     ) -> None:
         super().__init__()
@@ -430,7 +431,7 @@ class LocalAttention(nn.Module):
         if num_kv_heads is None:
             num_kv_heads = num_heads
 
-        dtype = get_compute_dtype()
+        dtype = compute_dtype or get_compute_dtype()
         attn_backend = get_attn_backend(
             head_size, dtype, supported_attention_backends=supported_attention_backends
         )

python/sglang/multimodal_gen/runtime/loader/utils.py

@@ -179,14 +179,20 @@ class skip_init_modules:
     def __enter__(self):
         # Save originals
         self._orig_reset = {}
-        for cls in (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d):
+        for cls in (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.Embedding):
             self._orig_reset[cls] = cls.reset_parameters
             cls.reset_parameters = lambda self: None  # skip init
+        from transformers.modeling_utils import PreTrainedModel
+
+        self._pretrained_model_cls = PreTrainedModel
+        self._orig_post_init = PreTrainedModel.post_init
+        PreTrainedModel.post_init = lambda self: None
 
     def __exit__(self, exc_type, exc_value, traceback):
         # restore originals
         for cls, orig in self._orig_reset.items():
             cls.reset_parameters = orig
+        self._pretrained_model_cls.post_init = self._orig_post_init
 
 
 def _normalize_component_type(module_type: str) -> str:

python/sglang/multimodal_gen/runtime/loader/weight_utils.py

@@ -9,7 +9,7 @@ import json
 import os
 import tempfile
 from collections import defaultdict
-from collections.abc import Generator, Iterable
+from collections.abc import Callable, Generator, Iterable
 from pathlib import Path
 
 import filelock
@@ -27,6 +27,7 @@ except ImportError:
 
 from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
@@ -183,12 +184,20 @@ def safetensors_weights_iterator(
     hf_weights_files: list[str],
     to_cpu: bool = True,
     use_runai_model_streamer: bool | None = None,
+    key_filter: Callable[[str], bool] | None = None,
+    clone_streamed_tensors: bool = True,
+    weight_load_plan: WeightLoadPlan | None = None,
 ) -> Generator[tuple[str, torch.Tensor], None, None]:
     """Iterate over the weights in the model safetensor files."""
     enable_tqdm = (
         not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
     )
-    device = "cpu" if to_cpu else str(get_local_torch_device())
+    if weight_load_plan is not None:
+        checkpoint_device = torch.device(weight_load_plan.checkpoint_load_device)
+        to_cpu = checkpoint_device.type == "cpu"
+        device = str(checkpoint_device)
+    else:
+        device = "cpu" if to_cpu else str(get_local_torch_device())
     if use_runai_model_streamer is None:
         use_runai_model_streamer = (
             HAS_RUNAI_MODEL_STREAMER and envs.SGLANG_USE_RUNAI_MODEL_STREAMER
@@ -233,13 +242,24 @@ def safetensors_weights_iterator(
     _raise_if_duplicate_safetensors_keys(hf_weights_files)
 
     if use_runai_model_streamer:
+        logger.info(
+            "Loading safetensors with Run:ai Model Streamer to %s",
+            "cpu" if to_cpu else device,
+        )
         with SafetensorsStreamer() as streamer:
-            streamer.stream_files(hf_weights_files)
+            if to_cpu:
+                streamer.stream_files(hf_weights_files)
+            else:
+                streamer.stream_files(hf_weights_files, device=device)
             for name, tensor in streamer.get_tensors():
+                if key_filter is not None and not key_filter(name):
+                    continue
                 if to_cpu:
                     yield name, tensor.clone().detach()
+                elif clone_streamed_tensors:
+                    yield name, tensor.clone().detach()
                 else:
-                    yield name, tensor.to(device)
+                    yield name, tensor
     else:
         for st_file in tqdm(
             hf_weights_files,
@@ -249,6 +269,8 @@ def safetensors_weights_iterator(
         ):
             with safe_open(st_file, framework="pt", device=device) as f:
                 for name in f.keys():  # noqa: SIM118
+                    if key_filter is not None and not key_filter(name):
+                        continue
                     param = f.get_tensor(name)
                     yield name, param
 

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -480,7 +480,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
             self._materialize_raw_frame_transport(output_batch, req)
         elif req.save_output and req.return_file_paths_only:
             self._materialize_file_path_transport(output_batch, save_output_paths)
-        elif req.return_frames:
+        elif getattr(req, "return_frames", False):
             self._materialize_frame_outputs_for_return(output_batch, req)
 
     def _materialize_raw_frame_transport(
@@ -518,7 +518,11 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
         self, output_batch: OutputBatch, req: Req
     ) -> None:
         """materialize the output from tensor to numpy frames for faster serialization"""
-        if self.rank != 0 or output_batch.output is None or not req.return_frames:
+        if (
+            self.rank != 0
+            or output_batch.output is None
+            or not getattr(req, "return_frames", False)
+        ):
             return
 
         if (
@@ -692,7 +696,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
             mismatched = [
                 field
                 for field in shared_output_fields
-                if getattr(req, field) != getattr(first_req, field)
+                if getattr(req, field, None) != getattr(first_req, field, None)
             ]
             if mismatched:
                 raise ValueError(

python/sglang/multimodal_gen/runtime/managers/scheduler.py

@@ -282,6 +282,9 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
             if len(reqs) == 1 or not allow_dynamic_batching:
                 return self.worker.execute_forward(reqs)
 
+            if self.server_args.pipeline_config.supports_native_grouped_requests():
+                return self._execute_generation_grouped(reqs)
+
             merged_req = self._try_merge_generation_reqs(reqs)
             if merged_req is None:
                 return self._execute_generation_sequential(reqs)
@@ -327,6 +330,48 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
                     error_msg=f"Dynamic batching failed: {e}",
                 )
 
+    def _execute_generation_grouped(self, reqs: List[Req]) -> List[OutputBatch]:
+        batch_size = len(reqs)
+        try:
+            output_batch = self.worker.execute_forward(reqs)
+            if output_batch.error:
+                logger.error(
+                    "Native grouped execution returned error. Returning per-request errors: %s",
+                    output_batch.error,
+                )
+                return self._build_dynamic_batch_error_outputs(
+                    reqs=reqs,
+                    error_msg=output_batch.error,
+                )
+
+            split_outputs = self._split_batched_output(output_batch, reqs)
+            if split_outputs is None:
+                logger.error(
+                    "Failed to split native grouped output cleanly. Returning per-request errors."
+                )
+                return self._build_dynamic_batch_error_outputs(
+                    reqs=reqs,
+                    error_msg="Native grouped execution failed: could not split output.",
+                )
+
+            logger.info(
+                "Processed native grouped batch of %d/%d request(s) with max_delay=%.2fms",
+                batch_size,
+                self._batching_max_size,
+                self._batching_delay_s * 1000.0,
+            )
+            return split_outputs
+        except Exception as e:
+            logger.error(
+                "Native grouped execution failed (%s). Returning per-request errors.",
+                e,
+                exc_info=True,
+            )
+            return self._build_dynamic_batch_error_outputs(
+                reqs=reqs,
+                error_msg=f"Native grouped execution failed: {e}",
+            )
+
     def _execute_generation_sequential(self, reqs: List[Req]) -> List[OutputBatch]:
         return [self.worker.execute_forward([req]) for req in reqs]
 
@@ -452,7 +497,10 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
             candidate_req.prompt, str
         ):
             return "prompt_type"
-        if base_req.image_path is not None or candidate_req.image_path is not None:
+        if (
+            getattr(base_req, "image_path", None) is not None
+            or getattr(candidate_req, "image_path", None) is not None
+        ):
             return "image_conditioning"
         if base_req.return_file_paths_only != candidate_req.return_file_paths_only:
             return "return_file_paths_only"
@@ -486,7 +534,10 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
         ):
             return False
 
-        if base_req.image_path is not None or candidate_req.image_path is not None:
+        if (
+            getattr(base_req, "image_path", None) is not None
+            or getattr(candidate_req, "image_path", None) is not None
+        ):
             return False
         if base_req.return_file_paths_only != candidate_req.return_file_paths_only:
             return False
@@ -722,8 +773,14 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
 
         outputs: list[OutputBatch] = []
         start = 0
-        for req, req_count in zip(reqs, per_req_counts):
+        for req_index, (req, req_count) in enumerate(zip(reqs, per_req_counts)):
             end = start + req_count
+            metrics = (
+                deepcopy(output_batch.metrics_list[req_index])
+                if output_batch.metrics_list is not None
+                and req_index < len(output_batch.metrics_list)
+                else deepcopy(output_batch.metrics)
+            )
             split = OutputBatch(
                 output=self._slice_batched_value(
                     output_batch.output, start, end, total_items
@@ -748,7 +805,7 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
                 output_file_paths=self._slice_batched_value(
                     output_batch.output_file_paths, start, end, total_items
                 ),
-                metrics=deepcopy(output_batch.metrics),
+                metrics=metrics,
                 noise_pred=self._slice_batched_value(
                     output_batch.noise_pred, start, end, total_items
                 ),

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

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from sglang.multimodal_gen.runtime.models.vlas.pi05_policy import Pi05PolicyModel
+
+__all__ = ["Pi05PolicyModel"]

python/sglang/multimodal_gen/runtime/pipelines/pi05.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import torch
+
+from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
+from sglang.multimodal_gen.configs.sample.pi05 import Pi05SamplingParams
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.models.vlas import Pi05PolicyModel
+from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
+    ComposedPipelineBase,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.pi05_preprocess import (
+    Pi05Preprocessor,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.vla import (
+    VLAActionDenoisingStage,
+    VLAActionPostprocessStage,
+    VLAObservationPreprocessStage,
+    VLAPrefixEncodingStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.vla.prefix_cache import VLAPrefixCacheManager
+
+logger = init_logger(__name__)
+
+
+class Pi05Pipeline(ComposedPipelineBase):
+    pipeline_name = "Pi05Pipeline"
+    pipeline_config_cls = Pi05PipelineConfig
+    sampling_params_cls = Pi05SamplingParams
+    _required_config_modules: list[str] = []
+
+    def validate_disagg_role(self, role: RoleType) -> None:
+        if role != RoleType.MONOLITHIC:
+            raise ValueError(
+                "Pi05Pipeline v1 supports same-process execution only. "
+                "Use prefix/action logical groups inside one worker; cross-node "
+                "multimodal_gen disaggregation is a v2 target."
+            )
+
+    def load_modules(
+        self,
+        server_args: ServerArgs,
+        loaded_modules: dict[str, torch.nn.Module] | None = None,
+    ) -> dict[str, torch.nn.Module]:
+        if loaded_modules is not None:
+            return loaded_modules
+
+        pipeline_config: Pi05PipelineConfig = server_args.pipeline_config
+        pipeline_config.offload_prefix_image_encoder = (
+            pipeline_config.offload_prefix_image_encoder
+            or bool(server_args.image_encoder_cpu_offload)
+        )
+        pipeline_config.offload_prefix_token_embedding = (
+            pipeline_config.offload_prefix_token_embedding
+            or bool(server_args.text_encoder_cpu_offload)
+        )
+        logger.info(
+            "Pi05 memory config: prefix_cache=%s/%s, action_cuda_graph=%s, "
+            "offload_image=%s, offload_image_after_embed=%s, "
+            "offload_tokens=%s, offload_language_layers=%s, "
+            "offload_language_after_prefix=%s/%s, "
+            "offload_action_after_denoise=%s, empty_cache_after_prefix=%s",
+            pipeline_config.enable_global_prefix_cache,
+            pipeline_config.prefix_cache_max_entries,
+            pipeline_config.enable_action_cuda_graph,
+            pipeline_config.offload_prefix_image_encoder,
+            pipeline_config.offload_prefix_image_encoder_after_embed,
+            pipeline_config.offload_prefix_token_embedding,
+            pipeline_config.offload_prefix_language_layers,
+            pipeline_config.offload_prefix_language_layers_after_prefix,
+            pipeline_config.offload_prefix_language_layer_count_after_prefix,
+            pipeline_config.offload_action_expert_after_denoise,
+            pipeline_config.empty_cache_after_prefix,
+        )
+        policy_model = Pi05PolicyModel.from_pretrained(
+            self.model_path,
+            pipeline_config,
+        )
+        if (
+            pipeline_config.prefix_parallel_strategy
+            == pipeline_config.action_parallel_strategy
+            == "tp"
+        ):
+            raise ValueError(
+                "VLA action expert should not share the prefix TP layout. "
+                "Use SP, Ulysses, Ring, DP, or monolithic fallback for the "
+                "action path."
+            )
+        return {
+            "policy_model": policy_model,
+        }
+
+    def initialize_pipeline(self, server_args: ServerArgs) -> None:
+        pipeline_config: Pi05PipelineConfig = server_args.pipeline_config
+        self.preprocessor = Pi05Preprocessor(pipeline_config)
+        self.prefix_cache = VLAPrefixCacheManager(
+            max_entries=pipeline_config.prefix_cache_max_entries
+        )
+
+    def create_pipeline_stages(self, server_args: ServerArgs):
+        self.add_stage(
+            VLAObservationPreprocessStage(self.preprocessor),
+            "pi05_preprocess",
+        )
+        self.add_stage(
+            VLAPrefixEncodingStage(
+                self.get_module("policy_model"),
+                self.prefix_cache,
+            ),
+            "pi05_prefix",
+        )
+        self.add_stage(
+            VLAActionDenoisingStage(self.get_module("policy_model")),
+            "pi05_action_denoise",
+        )
+        self.add_stage(
+            VLAActionPostprocessStage(),
+            "pi05_postprocess",
+        )
+
+
+EntryClass = Pi05Pipeline

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -979,7 +979,12 @@ class ComposedPipelineBase(ABC):
 
         # Execute each stage
         if not batch.is_warmup and not batch.suppress_logs:
-            logger.info(
+            stage_logger = (
+                logger.debug
+                if server_args.pipeline_config.task_type.is_action_gen()
+                else logger.info
+            )
+            stage_logger(
                 "Running pipeline stages: %s",
                 list(self._stage_name_mapping.keys()),
                 main_process_only=True,
@@ -1007,7 +1012,12 @@ class ComposedPipelineBase(ABC):
             )
 
         if not batches[0].is_warmup and not batches[0].suppress_logs:
-            logger.info(
+            stage_logger = (
+                logger.debug
+                if server_args.pipeline_config.task_type.is_action_gen()
+                else logger.info
+            )
+            stage_logger(
                 "Running grouped pipeline stages: %s",
                 list(self._stage_name_mapping.keys()),
                 main_process_only=True,

python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py

@@ -57,6 +57,10 @@ class PipelineExecutor(ABC):
         batch: Any,
         server_args: ServerArgs,
     ) -> None:
+        if isinstance(batch, list):
+            if not batch:
+                return
+            batch = batch[0]
         self.component_residency_manager.begin_request(stages, batch, server_args)
 
     def before_stage(

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -22,7 +22,10 @@ from typing import Any, Optional, Sequence, Union
 import PIL.Image
 import torch
 
-from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
 from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import (
     RolloutTrajectoryData,
 )
@@ -320,9 +323,11 @@ class Req:
     @property
     def resolution_key(self) -> str | None:
         """Return the batching config resolution key, e.g. "1024x1024"."""
-        if self.width is None or self.height is None:
+        width = getattr(self, "width", None)
+        height = getattr(self, "height", None)
+        if width is None or height is None:
             return None
-        return f"{int(self.width)}x{int(self.height)}"
+        return f"{int(width)}x{int(height)}"
 
     def set_as_warmup(self, warmup_steps: int = 1):
         self.is_warmup = True
@@ -339,6 +344,13 @@ class Req:
 
     def validate(self):
         """Initialize dependent fields after dataclass initialization."""
+        if getattr(self.sampling_params, "data_type", None) == DataType.ACTION:
+            self.do_classifier_free_guidance = False
+            if self.negative_prompt_embeds is None:
+                self.negative_prompt_embeds = []
+            self.metrics = RequestMetrics(request_id=self.request_id)
+            return
+
         # Prefer true_cfg_scale when it is explicitly provided.
         cfg_scale = (
             self.true_cfg_scale
@@ -360,6 +372,22 @@ class Req:
     def log(self, server_args: ServerArgs):
         if self.is_warmup or self.suppress_logs:
             return
+        if getattr(self.sampling_params, "data_type", None) == DataType.ACTION:
+            if not logger.isEnabledFor(logging.DEBUG):
+                return
+            logger.debug(
+                "VLA request: prompt=%s seed=%s steps=%s outputs=%s action=%sx%s "
+                "save_output=%s",
+                _sanitize_for_logging(self.prompt, key_hint="prompt"),
+                self.seed,
+                self.num_inference_steps,
+                self.num_outputs_per_prompt,
+                getattr(self, "action_horizon", None),
+                getattr(self, "action_dim", None),
+                self.save_output,
+            )
+            return
+
         # TODO: in some cases (e.g., TI2I), height and weight might be undecided at this moment
         if self.height:
             target_height = align_to(self.height, 16)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py

@@ -42,6 +42,12 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation im
     DMDTimestepPreparationStage,
     TimestepPreparationStage,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.vla import (
+    VLAActionDenoisingStage,
+    VLAActionPostprocessStage,
+    VLAObservationPreprocessStage,
+    VLAPrefixEncodingStage,
+)
 
 __all__ = [
     "PipelineStage",
@@ -60,4 +66,8 @@ __all__ = [
     "ImageEncodingStage",
     "ImageVAEEncodingStage",
     "TextEncodingStage",
+    "VLAObservationPreprocessStage",
+    "VLAPrefixEncodingStage",
+    "VLAActionDenoisingStage",
+    "VLAActionPostprocessStage",
 ]

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/pi05_preprocess.py

@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from transformers import AutoTokenizer
+
+from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
+from sglang.multimodal_gen.runtime.vla.observation import VLAObservationBatch
+
+
+def _tensor_from_image(value: Any) -> torch.Tensor:
+    if isinstance(value, torch.Tensor):
+        tensor = value.detach()
+        if tensor.ndim == 4:
+            if tensor.shape[0] != 1:
+                raise ValueError("Pi05 v1 expects one observation per request")
+            tensor = tensor[0]
+        if tensor.ndim != 3:
+            raise ValueError(f"Expected image tensor with 3 dims, got {tensor.shape}")
+        if tensor.shape[0] in (1, 3, 4):
+            tensor = tensor[:3]
+        elif tensor.shape[-1] in (1, 3, 4):
+            tensor = tensor[..., :3].permute(2, 0, 1)
+        else:
+            raise ValueError(
+                f"Could not infer image channels from shape {tensor.shape}"
+            )
+        is_integer = not tensor.is_floating_point()
+        tensor = tensor.to(dtype=torch.float32)
+        if is_integer or tensor.max() > 2.0:
+            tensor = tensor / 255.0
+        return tensor
+
+    if isinstance(value, Image.Image):
+        image = value.convert("RGB")
+        arr = np.asarray(image, dtype=np.float32) / 255.0
+        return torch.from_numpy(arr).permute(2, 0, 1)
+
+    if isinstance(value, (np.ndarray, list)):
+        arr = np.asarray(value)
+        if arr.ndim != 3:
+            raise ValueError(f"Expected HWC image array, got shape {arr.shape}")
+        tensor = torch.from_numpy(np.ascontiguousarray(arr))
+        if tensor.shape[0] in (1, 3, 4):
+            tensor = tensor[:3]
+        elif tensor.shape[-1] in (1, 3, 4):
+            tensor = tensor[..., :3].permute(2, 0, 1)
+        else:
+            raise ValueError(
+                f"Could not infer image channels from shape {tensor.shape}"
+            )
+        is_integer = not tensor.is_floating_point()
+        tensor = tensor.to(dtype=torch.float32)
+        if is_integer or tensor.max() > 2.0:
+            tensor = tensor / 255.0
+        return tensor
+
+    raise TypeError(f"Unsupported Pi05 image type: {type(value)}")
+
+
+def _resize_with_pad_image_tensor(
+    tensor: torch.Tensor, size: tuple[int, int]
+) -> torch.Tensor:
+    height, width = size
+    if tensor.shape[-2:] == (height, width):
+        return tensor
+    _, cur_height, cur_width = tensor.shape
+    ratio = max(cur_width / width, cur_height / height)
+    resized_height = int(cur_height / ratio)
+    resized_width = int(cur_width / ratio)
+    tensor = F.interpolate(
+        tensor[None],
+        size=(resized_height, resized_width),
+        mode="bilinear",
+        align_corners=False,
+    )[0]
+    pad_h0, rem_h = divmod(height - resized_height, 2)
+    pad_w0, rem_w = divmod(width - resized_width, 2)
+    return F.pad(
+        tensor,
+        (pad_w0, pad_w0 + rem_w, pad_h0, pad_h0 + rem_h),
+        mode="constant",
+        value=0.0,
+    )
+
+
+class Pi05Preprocessor:
+    def __init__(self, config: Pi05PipelineConfig):
+        self.config = config
+        self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name)
+        self.tokenizer.padding_side = "right"
+
+    def _tokenize(self, prompt: list[str], state: torch.Tensor | None):
+        if state is None:
+            state_for_prompt = torch.zeros(1, 0, dtype=torch.float32)
+        else:
+            state_for_prompt = state.detach().cpu().to(torch.float32)
+        bins = np.linspace(-1, 1, 256 + 1)[:-1]
+        state_np = state_for_prompt.numpy()
+        discretized = np.digitize(state_np, bins=bins) - 1
+        full_prompts = []
+        for idx, task in enumerate(prompt):
+            cleaned = task.strip().replace("_", " ").replace("\n", " ")
+            state_str = " ".join(map(str, discretized[idx]))
+            full_prompts.append(f"Task: {cleaned}, State: {state_str};\nAction: ")
+
+        encoded = self.tokenizer(
+            full_prompts,
+            max_length=self.config.max_token_len,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )
+        return encoded["input_ids"].to(torch.long), encoded["attention_mask"].to(
+            torch.bool
+        )
+
+    def __call__(self, raw_observation: dict[str, Any]) -> VLAObservationBatch:
+        prompt_value = raw_observation.get("prompt", "")
+        if isinstance(prompt_value, list):
+            prompt = [str(x) for x in prompt_value]
+        else:
+            prompt = [str(prompt_value)]
+        if len(prompt) != 1:
+            raise ValueError("Pi05 v1 expects one prompt per action request")
+
+        raw_images = raw_observation.get("images") or {}
+        image_masks_in = raw_observation.get("image_masks") or {}
+        camera_order = tuple(
+            raw_observation.get("camera_order") or self.config.image_keys
+        )
+
+        images: dict[str, torch.Tensor] = {}
+        image_masks: dict[str, torch.Tensor] = {}
+        for key in camera_order:
+            value = raw_images.get(key)
+            is_present = value is not None and bool(image_masks_in.get(key, True))
+            if is_present:
+                tensor = _tensor_from_image(value)
+                tensor = _resize_with_pad_image_tensor(tensor, self.config.image_size)
+                tensor = tensor * 2.0 - 1.0
+            else:
+                channels = 3
+                height, width = self.config.image_size
+                tensor = torch.ones(channels, height, width, dtype=torch.float32) * -1.0
+
+            images[key] = tensor.unsqueeze(0)
+            image_masks[key] = torch.tensor([is_present], dtype=torch.bool)
+
+        state = raw_observation.get("state")
+        state_tensor = None
+        if state is not None:
+            state_tensor = torch.as_tensor(state, dtype=torch.float32)
+            if state_tensor.ndim == 1:
+                state_tensor = state_tensor.unsqueeze(0)
+            if state_tensor.shape[0] != 1:
+                raise ValueError("Pi05 v1 expects one state vector per request")
+            if state_tensor.shape[-1] > self.config.state_dim:
+                raise ValueError(
+                    f"Pi05 state dim must be <= {self.config.state_dim}, "
+                    f"got {state_tensor.shape[-1]}"
+                )
+
+        noise = raw_observation.get("noise")
+        noise_tensor = None
+        if noise is not None:
+            noise_tensor = torch.as_tensor(noise, dtype=torch.float32)
+            if noise_tensor.ndim == 2:
+                noise_tensor = noise_tensor.unsqueeze(0)
+            expected = (1, self.config.action_horizon, self.config.action_dim)
+            if tuple(noise_tensor.shape) != expected:
+                raise ValueError(
+                    f"Pi05 noise must have shape {expected}, "
+                    f"got {tuple(noise_tensor.shape)}"
+                )
+
+        tokens = raw_observation.get("tokens")
+        if tokens is None:
+            tokens = raw_observation.get("tokenized_prompt")
+        token_masks = raw_observation.get("token_masks")
+        if token_masks is None:
+            token_masks = raw_observation.get("tokenized_prompt_mask")
+        if tokens is not None:
+            tokens_tensor = torch.as_tensor(tokens, dtype=torch.long)
+            if tokens_tensor.ndim == 1:
+                tokens_tensor = tokens_tensor.unsqueeze(0)
+            if token_masks is None:
+                token_masks_tensor = tokens_tensor != self.tokenizer.pad_token_id
+            else:
+                token_masks_tensor = torch.as_tensor(token_masks, dtype=torch.bool)
+                if token_masks_tensor.ndim == 1:
+                    token_masks_tensor = token_masks_tensor.unsqueeze(0)
+        else:
+            tokens_tensor, token_masks_tensor = self._tokenize(prompt, state_tensor)
+
+        return VLAObservationBatch(
+            prompt=prompt,
+            images=images,
+            image_masks=image_masks,
+            state=state_tensor,
+            noise=noise_tensor,
+            tokens=tokens_tensor,
+            token_masks=token_masks_tensor,
+            batch_size=1,
+            metadata={"camera_order": camera_order},
+        )

python/sglang/multimodal_gen/runtime/pipelines_core/stages/vla.py

@@ -0,0 +1,496 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import numpy as np
+import torch
+
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
+    OutputBatch,
+    Req,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.vla.observation import (
+    collate_vla_observation_batches,
+)
+from sglang.multimodal_gen.runtime.vla.parallel import (
+    broadcast_prefix_context,
+    broadcast_tensor_from_rank,
+    get_vla_split_group,
+)
+from sglang.multimodal_gen.runtime.vla.prefix_cache import (
+    PrefixContext,
+    VLAPrefixCacheManager,
+    slice_prefix_context,
+)
+
+
+def vla_state(batch: Req) -> dict[str, Any]:
+    """Per-request scratchpad shared by the VLA pipeline stages."""
+
+    return batch.extra["vla"]
+
+
+def vla_timings(batch: Req) -> dict[str, float]:
+    return vla_state(batch).setdefault("timings", {})
+
+
+def vla_options(batch: Req) -> dict[str, Any]:
+    return vla_state(batch).get("options") or {}
+
+
+def materialize_vla_action_batch(
+    actions: Any,
+    action_dim: int,
+    output_format: str,
+) -> Any:
+    output_format = output_format.lower()
+    if isinstance(actions, torch.Tensor):
+        actions_out = actions[..., :action_dim].detach().float().cpu().numpy()
+        if output_format != "numpy":
+            actions_out = actions_out.tolist()
+    elif isinstance(actions, np.ndarray):
+        actions_out = actions[..., :action_dim].astype(np.float32, copy=False)
+        if output_format != "numpy":
+            actions_out = actions_out.tolist()
+    else:
+        actions_out = actions
+    if not isinstance(actions_out, list):
+        return actions_out
+    if not actions_out:
+        return []
+    first = actions_out[0]
+    if isinstance(first, list) and first and isinstance(first[0], list):
+        actions_out = [[step[:action_dim] for step in sample] for sample in actions_out]
+    else:
+        actions_out = [[step[:action_dim] for step in actions_out]]
+    if output_format == "numpy":
+        return np.asarray(actions_out, dtype=np.float32)
+    return actions_out
+
+
+def synchronize_vla_action_tensor(actions: torch.Tensor | None) -> None:
+    if actions is not None and actions.device.type == "cuda":
+        torch.cuda.synchronize(actions.device)
+
+
+def _effective_prefix_cache_enabled(
+    batch: Req,
+    server_args: ServerArgs,
+) -> bool:
+    options = vla_options(batch)
+    return bool(options.get("enable_prefix_cache", True)) and bool(
+        server_args.pipeline_config.enable_global_prefix_cache
+    )
+
+
+def _grouped_fingerprint(
+    batch: Req,
+    server_args: ServerArgs,
+) -> tuple[Any, ...]:
+    if (
+        batch.is_warmup
+        or get_vla_split_group() is not None
+        or _effective_prefix_cache_enabled(batch, server_args)
+        or batch.generator is not None
+    ):
+        return ("single", id(batch))
+
+    observation = vla_state(batch).get("observation_batch")
+    camera_order = tuple(observation.metadata.get("camera_order", ()))
+    image_shapes = tuple(
+        (
+            name,
+            tuple(observation.images[name].shape),
+            bool(observation.image_masks[name].item()),
+        )
+        for name in camera_order
+    )
+    return (
+        "grouped",
+        camera_order,
+        image_shapes,
+        None if observation.state is None else tuple(observation.state.shape),
+        None if observation.noise is None else tuple(observation.noise.shape),
+        tuple(observation.tokens.shape),
+        tuple(observation.token_masks.shape),
+        batch.action_horizon,
+        batch.action_dim,
+        batch.num_inference_steps,
+    )
+
+
+class VLAObservationPreprocessStage(PipelineStage):
+    def __init__(self, preprocessor: Any):
+        super().__init__()
+        self.preprocessor = preprocessor
+
+    @property
+    def role_affinity(self) -> RoleType:
+        return RoleType.ENCODER
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        start = time.perf_counter()
+        state = vla_state(batch)
+        raw_observation = dict(state.get("observation") or {})
+        raw_observation.setdefault("prompt", batch.prompt)
+        observation = self.preprocessor(raw_observation)
+        state["observation_batch"] = observation
+        vla_timings(batch)["preprocess_ms"] = (time.perf_counter() - start) * 1000
+        return batch
+
+
+class VLAPrefixEncodingStage(PipelineStage):
+    def __init__(
+        self,
+        policy_model: Any,
+        prefix_cache: VLAPrefixCacheManager,
+    ):
+        super().__init__()
+        self.policy_model = policy_model
+        self.prefix_cache = prefix_cache
+
+    @property
+    def role_affinity(self) -> RoleType:
+        return RoleType.ENCODER
+
+    def run_grouped_requests(
+        self,
+        batches: list[Req],
+        server_args: ServerArgs,
+    ) -> list[Req]:
+        results: list[Req | None] = [None] * len(batches)
+        for fingerprint, group in self._group_requests_by_fingerprint(
+            batches,
+            lambda batch: _grouped_fingerprint(batch, server_args),
+        ):
+            group_batches = [batch for _, batch in group]
+            if len(group_batches) == 1 or fingerprint[0] == "single":
+                for index, batch in group:
+                    results[index] = self(batch, server_args)
+                continue
+
+            prefix_start = time.perf_counter()
+            observations = [
+                vla_state(batch)["observation_batch"] for batch in group_batches
+            ]
+            grouped_observation = collate_vla_observation_batches(observations)
+            prefix_context = self.policy_model.encode_prefix(grouped_observation)
+            prefix_ms = (time.perf_counter() - prefix_start) * 1000
+
+            for offset, (index, batch) in enumerate(group):
+                state = vla_state(batch)
+                state["observation_group"] = grouped_observation
+                state["prefix_context_group"] = prefix_context
+                state["prefix_context"] = slice_prefix_context(
+                    prefix_context,
+                    offset,
+                )
+                state["cache"] = {
+                    "hit": False,
+                    "scope": "request",
+                    "prefix_len": prefix_context.prefix_len,
+                    "grouped": True,
+                    "batch_size": len(group_batches),
+                }
+                timings = vla_timings(batch)
+                timings["cache_lookup_ms"] = 0.0
+                timings["prefix_ms"] = prefix_ms
+                results[index] = batch
+
+            if (
+                server_args.pipeline_config.empty_cache_after_prefix
+                and torch.cuda.is_available()
+            ):
+                torch.cuda.empty_cache()
+
+        return [result for result in results if result is not None]
+
+    def _recv_prefix_result(self, batch: Req, split: Any) -> Req:
+        state = vla_state(batch)
+        state["prefix_context"] = broadcast_prefix_context(
+            None,
+            split,
+            src=split.prefix_root,
+        )
+        state["cache"] = split.broadcast_object_from_rank(
+            None,
+            src=split.prefix_root,
+        )
+        timings = split.broadcast_object_from_rank(None, src=split.prefix_root)
+        vla_timings(batch).update(timings)
+        return batch
+
+    def _send_prefix_result(
+        self,
+        batch: Req,
+        split: Any,
+        prefix_context: Any,
+    ) -> None:
+        broadcast_prefix_context(
+            prefix_context,
+            split,
+            src=split.prefix_root,
+        )
+        split.broadcast_object_from_rank(
+            vla_state(batch)["cache"],
+            src=split.prefix_root,
+        )
+        split.broadcast_object_from_rank(
+            vla_timings(batch),
+            src=split.prefix_root,
+        )
+
+    def get_cached_context(
+        self, batch: Req, server_args: ServerArgs, observation: Any
+    ) -> tuple[str, PrefixContext]:
+        """try querying the cache for PrefixContext with prefix cache key built from observations and other keys"""
+        cache_enabled = _effective_prefix_cache_enabled(batch, server_args)
+        if cache_enabled:
+            cache_key = self.policy_model.build_prefix_cache_key(observation)
+            cached_context = self.prefix_cache.get(cache_key)
+        else:
+            cache_key = None
+            cached_context = None
+        return cache_key, cached_context
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        state = vla_state(batch)
+        if batch.is_warmup:
+            state["prefix_context"] = None
+            state["cache"] = {"hit": False, "warmup": True}
+            return batch
+
+        split = get_vla_split_group()
+        if split is not None and not split.is_prefix_rank:
+            return self._recv_prefix_result(batch, split)
+
+        observation = state["observation_batch"]
+        cache_start = time.perf_counter()
+        cache_enabled = _effective_prefix_cache_enabled(batch, server_args)
+
+        # 1. try querying the per-request LRU prefix kv cache
+        cache_key, cached_context = self.get_cached_context(
+            batch, server_args, observation
+        )
+
+        vla_timings(batch)["cache_lookup_ms"] = (
+            time.perf_counter() - cache_start
+        ) * 1000
+
+        # 2. prepare VLAState
+        if cached_context is not None:
+            state["prefix_context"] = cached_context
+            state["cache"] = {
+                "hit": True,
+                "scope": "global",
+                "mode": "exact",
+                "prefix_len": cached_context.prefix_len,
+            }
+            if split is not None:
+                self._send_prefix_result(batch, split, cached_context)
+            return batch
+
+        prefix_start = time.perf_counter()
+
+        # 3. run encoding
+        prefix_context = self.policy_model.encode_prefix(observation)
+        if cache_key is not None:
+            prefix_context.cache_key_digest = cache_key
+
+        vla_timings(batch)["prefix_ms"] = (time.perf_counter() - prefix_start) * 1000
+        state["prefix_context"] = prefix_context
+        state["cache"] = {
+            "hit": False,
+            "scope": "global" if cache_enabled else "request",
+            "mode": "exact" if cache_enabled else "disabled",
+            "prefix_len": prefix_context.prefix_len,
+        }
+
+        # 4. update prefix kv cache
+        if cache_key is not None:
+            self.prefix_cache.put(cache_key, prefix_context)
+        if split is not None:
+            self._send_prefix_result(batch, split, prefix_context)
+        if (
+            server_args.pipeline_config.empty_cache_after_prefix
+            and torch.cuda.is_available()
+        ):
+            torch.cuda.empty_cache()
+        return batch
+
+
+class VLAActionDenoisingStage(PipelineStage):
+    def __init__(self, policy_model: Any):
+        super().__init__()
+        self.policy_model = policy_model
+
+    @property
+    def role_affinity(self) -> RoleType:
+        return RoleType.DENOISER
+
+    def run_grouped_requests(
+        self,
+        batches: list[Req],
+        server_args: ServerArgs,
+    ) -> list[Req]:
+        results: list[Req | None] = [None] * len(batches)
+
+        def action_fingerprint(batch: Req) -> tuple[Any, ...]:
+            prefix_context = vla_state(batch).get("prefix_context_group")
+            if prefix_context is None:
+                return ("single", id(batch))
+            return (
+                "grouped",
+                id(prefix_context),
+                batch.num_inference_steps,
+                str(vla_options(batch).get("output_format") or "list"),
+            )
+
+        for _, group in self._group_requests_by_fingerprint(
+            batches,
+            action_fingerprint,
+        ):
+            group_batches = [batch for _, batch in group]
+            prefix_context = vla_state(group_batches[0]).get("prefix_context_group")
+            if len(group_batches) == 1 or prefix_context is None:
+                for index, batch in group:
+                    results[index] = self(batch, server_args)
+                continue
+
+            start = time.perf_counter()
+            options = vla_options(group_batches[0])
+            observation = vla_state(group_batches[0])["observation_group"]
+            actions = self.policy_model.sample_actions(
+                observation,
+                prefix_context,
+                noise=observation.noise,
+                num_steps=group_batches[0].num_inference_steps,
+                use_cuda_graph=bool(options.get("enable_cuda_graph", True)),
+                generator=None,
+            )
+            synchronize_vla_action_tensor(actions)
+            actions_out = materialize_vla_action_batch(
+                actions,
+                server_args.pipeline_config.output_action_dim,
+                str(options.get("output_format") or "list"),
+            )
+            action_ms = (time.perf_counter() - start) * 1000
+            parallel_info = self.policy_model.action_parallel_info(prefix_context)
+
+            for offset, (index, batch) in enumerate(group):
+                state = vla_state(batch)
+                vla_timings(batch)["action_denoise_ms"] = action_ms
+                state["parallel"] = parallel_info
+                state["actions_output"] = actions_out[offset]
+                results[index] = batch
+
+        return [result for result in results if result is not None]
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        start = time.perf_counter()
+        state = vla_state(batch)
+        observation = state.get("observation_batch")
+        split = get_vla_split_group()
+        prefix_context = state.get("prefix_context")
+        should_run_action = (
+            split is None or self.policy_model.should_run_action_denoise(prefix_context)
+        )
+        parallel_info = self.policy_model.action_parallel_info(prefix_context)
+        if batch.is_warmup:
+            actions = (
+                self.policy_model.warmup_actions(batch_size=1)
+                if should_run_action
+                else None
+            )
+        elif should_run_action:
+            # broadcast PrefixContext from action root rank to action ranks
+            options = vla_options(batch)
+            noise = observation.noise if observation is not None else None
+            actions = self.policy_model.sample_actions(
+                observation,
+                prefix_context,
+                noise=noise,
+                num_steps=batch.num_inference_steps,
+                use_cuda_graph=bool(options.get("enable_cuda_graph", True)),
+                generator=batch.generator,
+            )
+            synchronize_vla_action_tensor(actions)
+        else:
+            actions = None
+
+        if split is not None:
+            if should_run_action:
+                vla_timings(batch)["action_denoise_ms"] = (
+                    time.perf_counter() - start
+                ) * 1000
+            actions = broadcast_tensor_from_rank(
+                actions,
+                split,
+                src=split.action_root,
+                device=self.policy_model.device,
+            )
+            timings = split.broadcast_object_from_rank(
+                vla_timings(batch) if should_run_action else None,
+                src=split.action_root,
+            )
+            vla_timings(batch).update(timings)
+            parallel_info = split.broadcast_object_from_rank(
+                parallel_info if should_run_action else None,
+                src=split.action_root,
+            )
+        else:
+            vla_timings(batch)["action_denoise_ms"] = (
+                time.perf_counter() - start
+            ) * 1000
+        state["parallel"] = parallel_info
+        state["actions"] = actions
+        return batch
+
+
+class VLAActionPostprocessStage(PipelineStage):
+    @property
+    def role_affinity(self) -> RoleType:
+        return RoleType.DENOISER
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch:
+        start = time.perf_counter()
+        state = vla_state(batch)
+        action_dim = server_args.pipeline_config.output_action_dim
+        options = vla_options(batch)
+        actions_out = state.get("actions_output")
+        if actions_out is None:
+            action_batch = materialize_vla_action_batch(
+                state["actions"],
+                action_dim,
+                str(options.get("output_format") or "list"),
+            )
+            actions_out = (
+                action_batch[0] if isinstance(action_batch, list) else action_batch
+            )
+            if isinstance(action_batch, np.ndarray):
+                actions_out = action_batch[0]
+
+        payload = {
+            "request_id": batch.request_id,
+            "actions": actions_out,
+        }
+        payload["parameters"] = {"num_inference_steps": batch.num_inference_steps}
+        if options.get("return_timing", True):
+            timings = dict(vla_timings(batch))
+            timings["postprocess_ms"] = (time.perf_counter() - start) * 1000
+            payload["timings"] = timings
+        if not batch.is_warmup:
+            payload["cache"] = state.get("cache", {})
+        if state.get("parallel") is not None:
+            payload["parallel"] = state["parallel"]
+
+        return OutputBatch(
+            output=[payload],
+            metrics=batch.metrics,
+        )

python/sglang/multimodal_gen/runtime/server_args/auto_tune.py

@@ -398,6 +398,8 @@ class ServerArgsAutoTuner:
 
     def _default_layerwise_components_for_unset_placement(self) -> list[str]:
         args = self.server_args
+        if args.pipeline_config.task_type.is_action_gen():
+            return []
         if (
             args.is_arg_explicitly_set("layerwise_offload_components")
             or args.dit_layerwise_offload is True

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -615,6 +615,17 @@ class ServerArgs(DisaggServerArgsMixin):
             # CPU platform does not need offload
             return
 
+        if self.pipeline_config.task_type.is_action_gen():
+            if self.dit_cpu_offload is None:
+                self.dit_cpu_offload = False
+            if self.text_encoder_cpu_offload is None:
+                self.text_encoder_cpu_offload = False
+            if self.image_encoder_cpu_offload is None:
+                self.image_encoder_cpu_offload = False
+            if self.vae_cpu_offload is None:
+                self.vae_cpu_offload = False
+            return
+
         # TODO: to be handled by each platform
         if current_platform.get_device_total_memory() / BYTES_PER_GB < 30:
             logger.info(
@@ -1106,6 +1117,20 @@ class ServerArgs(DisaggServerArgsMixin):
             self.use_fsdp_inference = False
             self.dit_layerwise_offload = False
             self.layerwise_offload_components = None
+            if (
+                self.dit_cpu_offload
+                or self.text_encoder_cpu_offload
+                or self.image_encoder_cpu_offload
+                or self.vae_cpu_offload
+            ):
+                logger.warning(
+                    "Disabling component CPU offload on MPS because CPU-to-MPS "
+                    "module relocation can produce invalid diffusion outputs."
+                )
+            self.dit_cpu_offload = False
+            self.text_encoder_cpu_offload = False
+            self.image_encoder_cpu_offload = False
+            self.vae_cpu_offload = False
 
     def is_arg_explicitly_set(self, arg_name: str) -> bool:
         return arg_name in self._explicit_arg_names
@@ -1283,12 +1308,14 @@ class ServerArgs(DisaggServerArgsMixin):
             ),
         )
         parser.add_argument(
+            "--pipeline",
             "--pipeline-class-name",
+            dest="pipeline_class_name",
             type=str,
             default=ServerArgs.pipeline_class_name,
             help=(
-                "Override pipeline class selection from model_index.json. "
-                "Must match a registered pipeline_name."
+                "Advanced override for pipeline class selection from the model registry "
+                "or model_index.json. Must match a registered pipeline_name."
             ),
         )
         # attention
@@ -2173,7 +2200,7 @@ class ServerArgs(DisaggServerArgsMixin):
 
         # Create a set of argument names that were present on the command line.
         # This handles both styles: '--arg=value' and '--arg value'.
-        provided_arg_names = set()
+        provided_arg_names = set(getattr(args, "_sglang_explicit_arg_names", ()))
         for arg in raw_argv:
             if arg.startswith("--"):
                 # For '--arg=value', this gets 'arg'; for '--arg', this also gets 'arg'.
@@ -2192,6 +2219,8 @@ class ServerArgs(DisaggServerArgsMixin):
 
         # Populate provided_args if the argument from the namespace was on the command line.
         for k, v in vars(args).items():
+            if k.startswith("_sglang_"):
+                continue
             if k in provided_arg_names:
                 provided_args[k] = v
 

python/sglang/multimodal_gen/runtime/server_warmup.py

@@ -17,6 +17,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.warmup_request_builder import (
     build_warmup_reqs,
     should_include_warmup_image,
+    supports_synthetic_warmup,
 )
 
 logger = init_logger(__name__)
@@ -85,13 +86,19 @@ def is_realtime_serving(server_args: ServerArgs) -> bool:
 
 
 def should_run_synthetic_server_warmup(server_args: ServerArgs) -> bool:
-    return should_run_server_warmup(server_args) and not is_realtime_serving(
-        server_args
+    return (
+        should_run_server_warmup(server_args)
+        and supports_synthetic_warmup(server_args)
+        and not is_realtime_serving(server_args)
     )
 
 
 def should_run_explicit_client_warmup(server_args: ServerArgs) -> bool:
-    return server_args.warmup and server_args.warmup_resolutions is not None
+    return (
+        server_args.warmup
+        and server_args.warmup_resolutions is not None
+        and supports_synthetic_warmup(server_args)
+    )
 
 
 def format_warmup_req(req_or_group: Any) -> str:
@@ -102,9 +109,12 @@ def format_warmup_req(req_or_group: Any) -> str:
     if req is None:
         return prefix
 
-    shape = f"{req.width}x{req.height}"
-    if req.num_frames is not None and req.num_frames > 1:
-        shape += f"x{req.num_frames}f"
+    width = getattr(req, "width", None)
+    height = getattr(req, "height", None)
+    shape = "action" if width is None or height is None else f"{width}x{height}"
+    num_frames = getattr(req, "num_frames", None)
+    if num_frames is not None and num_frames > 1:
+        shape += f"x{num_frames}f"
 
     default_steps = req.extra.get("cache_dit_num_inference_steps")
     if default_steps is not None and default_steps != req.num_inference_steps:

python/sglang/multimodal_gen/runtime/vla/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared VLA runtime contracts and execution infrastructure."""

python/sglang/multimodal_gen/runtime/vla/denoise_cuda_graph.py

@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable
+
+import torch
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.vla.prefix_cache import (
+    PrefixContext,
+    VLADensePrefixCache,
+)
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_memory_pool,
+)
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class VLADenoiseGraphSignature:
+    batch_size: int
+    prefix_len: int
+    action_horizon: int
+    action_dim: int
+    dtype: str
+    parallel_layout: str
+
+
+@dataclass
+class _CapturedDenoiseGraph:
+    graph: torch.cuda.CUDAGraph
+    static_prefix_context: PrefixContext
+    static_x_t: torch.Tensor
+    static_timestep: torch.Tensor
+    static_output: torch.Tensor
+    current_context_id: int | None = None
+    current_context_digest: str | None = None
+
+
+def _clone_past_key_values(past_key_values: Any) -> Any:
+    return VLADensePrefixCache(
+        tuple(
+            (keys.detach().clone(), values.detach().clone(), sliding_window)
+            for keys, values, sliding_window in past_key_values
+        )
+    )
+
+
+def _copy_past_key_values_(dst: Any, src: Any) -> None:
+    for (dst_keys, dst_values, _), (src_keys, src_values, _) in zip(
+        dst, src, strict=True
+    ):
+        dst_keys.copy_(src_keys)
+        dst_values.copy_(src_values)
+
+
+def _clone_prefix_context(prefix_context: PrefixContext) -> PrefixContext:
+    return PrefixContext(
+        past_key_values=_clone_past_key_values(prefix_context.past_key_values),
+        prefix_pad_masks=prefix_context.prefix_pad_masks.detach().clone(),
+        prefix_len=prefix_context.prefix_len,
+        layout=dict(prefix_context.layout),
+        cache_key_digest=prefix_context.cache_key_digest,
+    )
+
+
+def _copy_prefix_context_(dst: PrefixContext, src: PrefixContext) -> None:
+    dst.prefix_pad_masks.copy_(src.prefix_pad_masks)
+    _copy_past_key_values_(dst.past_key_values, src.past_key_values)
+    dst.cache_key_digest = src.cache_key_digest
+
+
+class VLADenoiseGraphRunner:
+    """Full CUDA graph runner for one VLA action-denoise step.
+
+    Each signature owns fixed input and output buffers. This does not use
+    diffusion BCG and does not capture prefix encoding or token decode.
+    """
+
+    def __init__(self, enabled: bool = True):
+        self.enabled = enabled
+        self._captured: dict[VLADenoiseGraphSignature, _CapturedDenoiseGraph] = {}
+        self._disabled_signatures: set[VLADenoiseGraphSignature] = set()
+        self._capture_stream: torch.cuda.Stream | None = None
+        self._graph_pool: Any = None
+
+    def _sync_context_if_needed(
+        self,
+        captured: _CapturedDenoiseGraph,
+        prefix_context: PrefixContext,
+    ) -> None:
+        context_id = id(prefix_context.past_key_values)
+        context_digest = prefix_context.cache_key_digest
+        if (
+            context_digest is not None
+            and captured.current_context_digest == context_digest
+        ):
+            captured.current_context_id = context_id
+            return
+        if captured.current_context_id == context_id:
+            return
+        _copy_prefix_context_(captured.static_prefix_context, prefix_context)
+        captured.current_context_id = context_id
+        captured.current_context_digest = context_digest
+
+    def _capture(
+        self,
+        signature: VLADenoiseGraphSignature,
+        step_fn: Callable[..., torch.Tensor],
+        prefix_context: PrefixContext,
+        x_t: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> _CapturedDenoiseGraph:
+        static_prefix_context = _clone_prefix_context(prefix_context)
+        static_x_t = x_t.detach().clone()
+        static_timestep = timestep.detach().clone()
+
+        device_module = torch.get_device_module(x_t.device)
+        if self._capture_stream is None:
+            self._capture_stream = device_module.Stream(device=x_t.device)
+        if self._graph_pool is None:
+            self._graph_pool = get_or_create_global_graph_memory_pool(device_module)
+            set_graph_pool_id(self._graph_pool)
+
+        # warm up lazy kernels and workspaces before capture
+        device_module.synchronize()
+        with device_module.stream(self._capture_stream), torch.inference_mode():
+            step_fn(
+                static_prefix_context,
+                static_x_t,
+                static_timestep,
+            )
+        self._capture_stream.synchronize()
+
+        graph = torch.cuda.CUDAGraph()
+        with (
+            device_module.graph(
+                cuda_graph=graph,
+                pool=self._graph_pool,
+                stream=self._capture_stream,
+            ),
+            torch.inference_mode(),
+        ):
+            static_output = step_fn(
+                static_prefix_context,
+                static_x_t,
+                static_timestep,
+            )
+        self._capture_stream.synchronize()
+
+        captured = _CapturedDenoiseGraph(
+            graph=graph,
+            static_prefix_context=static_prefix_context,
+            static_x_t=static_x_t,
+            static_timestep=static_timestep,
+            static_output=static_output,
+            current_context_id=id(prefix_context.past_key_values),
+            current_context_digest=prefix_context.cache_key_digest,
+        )
+        self._captured[signature] = captured
+        logger.info(
+            "Captured VLA denoise CUDA graph: batch=%d prefix=%d action=%dx%d "
+            "dtype=%s",
+            signature.batch_size,
+            signature.prefix_len,
+            signature.action_horizon,
+            signature.action_dim,
+            signature.dtype,
+        )
+        return captured
+
+    def capture_or_run(
+        self,
+        signature: VLADenoiseGraphSignature,
+        step_fn: Callable[..., torch.Tensor],
+        prefix_context: PrefixContext,
+        x_t: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        if not self.enabled or signature in self._disabled_signatures:
+            return step_fn(prefix_context, x_t, timestep)
+
+        if x_t.device.type != "cuda":
+            return step_fn(prefix_context, x_t, timestep)
+
+        captured = self._captured.get(signature)
+        try:
+            if captured is None:
+                captured = self._capture(
+                    signature, step_fn, prefix_context, x_t, timestep
+                )
+                captured.graph.replay()
+            else:
+                self._sync_context_if_needed(captured, prefix_context)
+                captured.static_x_t.copy_(x_t)
+                captured.static_timestep.copy_(timestep)
+                captured.graph.replay()
+            return captured.static_output
+        except Exception:
+            self._disabled_signatures.add(signature)
+            self._captured.pop(signature, None)
+            logger.warning(
+                "VLA denoise CUDA graph disabled for signature %s",
+                signature,
+                exc_info=True,
+            )
+            return step_fn(prefix_context, x_t, timestep)

python/sglang/multimodal_gen/runtime/vla/observation.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from sglang.srt.managers.mm_utils import tensor_hash
+
+
+@dataclass
+class VLAObservationBatch:
+    prompt: list[str]
+    images: dict[str, torch.Tensor]
+    image_masks: dict[str, torch.Tensor]
+    state: torch.Tensor | None
+    noise: torch.Tensor | None
+    tokens: torch.Tensor
+    token_masks: torch.Tensor
+    batch_size: int
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+def tensor_fingerprint(tensor: torch.Tensor) -> str:
+    """Hash tensor content with SRT's CPU/CUDA implementation."""
+
+    shape = ",".join(str(dim) for dim in tensor.shape)
+    return f"{tensor.dtype}:{shape}:{tensor_hash(tensor):016x}"
+
+
+def collate_vla_observation_batches(
+    observations: list[VLAObservationBatch],
+) -> VLAObservationBatch:
+    first = observations[0]
+    camera_order = tuple(first.metadata.get("camera_order", ()))
+    images = {
+        name: torch.cat([obs.images[name] for obs in observations], dim=0)
+        for name in camera_order
+    }
+    image_masks = {
+        name: torch.cat([obs.image_masks[name] for obs in observations], dim=0)
+        for name in camera_order
+    }
+    states = [obs.state for obs in observations]
+    noises = [obs.noise for obs in observations]
+    if any(item is None for item in states) and not all(
+        item is None for item in states
+    ):
+        raise ValueError("Cannot collate mixed VLA state presence")
+    if any(item is None for item in noises) and not all(
+        item is None for item in noises
+    ):
+        raise ValueError("Cannot collate mixed VLA noise presence")
+    state = (
+        None
+        if states[0] is None
+        else torch.cat([item for item in states if item is not None], dim=0)
+    )
+    noise = (
+        None
+        if noises[0] is None
+        else torch.cat([item for item in noises if item is not None], dim=0)
+    )
+    return VLAObservationBatch(
+        prompt=[prompt for obs in observations for prompt in obs.prompt],
+        images=images,
+        image_masks=image_masks,
+        state=state,
+        noise=noise,
+        tokens=torch.cat([obs.tokens for obs in observations], dim=0),
+        token_masks=torch.cat([obs.token_masks for obs in observations], dim=0),
+        batch_size=len(observations),
+        metadata={"camera_order": camera_order},
+    )

python/sglang/multimodal_gen/runtime/vla/parallel.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+
+from sglang.multimodal_gen.runtime.distributed import (
+    get_sp_group,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.distributed.group_coordinator import GroupCoordinator
+from sglang.multimodal_gen.runtime.distributed.parallel_state import get_world_rank
+from sglang.multimodal_gen.runtime.vla.prefix_cache import (
+    PrefixContext,
+    VLADensePrefixCache,
+)
+
+
+@dataclass(frozen=True)
+class VLASplitGroup:
+    """Runtime view for VLA prefix/action split execution.
+
+    This reuses the existing SP group as the coordination group. It is not a
+    separate parallel topology:
+    1. `prefix_root` computes/fetches PrefixContext and broadcasts it once.
+    2. `action_root` owns fallback action denoise and initial noise broadcast.
+    3. `action_ranks` may all participate in action SP when the policy allows it.
+
+    All rank fields are global ranks; GroupCoordinator APIs take group-local
+    ranks, so call `group_rank_for` before collective helpers.
+    """
+
+    group: GroupCoordinator
+    prefix_root: int
+    action_root: int
+    action_ranks: tuple[int, ...]
+    rank: int
+
+    @property
+    def is_prefix_rank(self) -> bool:
+        return self.rank == self.prefix_root
+
+    @property
+    def is_action_rank(self) -> bool:
+        return self.rank in self.action_ranks
+
+    @property
+    def uses_action_sp(self) -> bool:
+        return len(self.action_ranks) > 1
+
+    def group_rank_for(self, global_rank: int) -> int:
+        return self.group.ranks.index(global_rank)
+
+    def broadcast_object_from_rank(self, obj, *, src: int):
+        return self.group.broadcast_object(
+            obj if self.rank == src else None,
+            src=self.group_rank_for(src),
+        )
+
+
+def get_vla_split_group() -> VLASplitGroup | None:
+    if not dist.is_available() or not dist.is_initialized():
+        return None
+    if not model_parallel_is_initialized():
+        return None
+    group = get_sp_group()
+    if group.world_size <= 1:
+        return None
+    # v1 maps the split view onto SP: first rank does prefix encode, last rank
+    # is the action fallback/root, and all SP ranks are eligible action ranks.
+    return VLASplitGroup(
+        group=group,
+        prefix_root=group.ranks[0],
+        action_root=group.ranks[-1],
+        action_ranks=tuple(group.ranks),
+        rank=get_world_rank(),
+    )
+
+
+def broadcast_tensor_from_rank(
+    tensor: torch.Tensor | None,
+    split: VLASplitGroup,
+    *,
+    src: int,
+    device: torch.device,
+) -> torch.Tensor | None:
+    payload = (
+        {"is_none": tensor is None, "tensor": tensor} if split.rank == src else None
+    )
+    payload = split.group.broadcast_tensor_dict(
+        payload,
+        src=split.group_rank_for(src),
+    )
+    if payload["is_none"]:
+        return None
+    output = payload["tensor"]
+    if output.device != device:
+        output = output.to(device)
+    return output
+
+
+def broadcast_prefix_context(
+    context: PrefixContext | None,
+    split: VLASplitGroup,
+    *,
+    src: int,
+) -> PrefixContext | None:
+    if split.rank == src and context is None:
+        payload = {"is_none": True}
+    elif split.rank == src:
+        prefix_pad_masks = context.prefix_pad_masks
+        prefix_pad_masks_is_bool = prefix_pad_masks.dtype == torch.bool
+        if prefix_pad_masks_is_bool:
+            prefix_pad_masks = prefix_pad_masks.to(torch.uint8)
+        payload = {
+            "is_none": False,
+            "prefix_pad_masks": prefix_pad_masks,
+            "prefix_pad_masks_is_bool": prefix_pad_masks_is_bool,
+            "prefix_len": context.prefix_len,
+            "layout": dict(context.layout),
+            "cache_key_digest": context.cache_key_digest,
+            "num_layers": len(context.past_key_values),
+        }
+        for i, (keys, values, sliding_window) in enumerate(context.past_key_values):
+            payload[f"layer_{i}_keys"] = keys
+            payload[f"layer_{i}_values"] = values
+            payload[f"layer_{i}_sliding_window"] = sliding_window
+    else:
+        payload = None
+
+    payload = split.group.broadcast_tensor_dict(
+        payload,
+        src=split.group_rank_for(src),
+    )
+    if payload["is_none"]:
+        return None
+
+    kv_layers = []
+    for i in range(int(payload["num_layers"])):
+        kv_layers.append(
+            (
+                payload[f"layer_{i}_keys"],
+                payload[f"layer_{i}_values"],
+                payload[f"layer_{i}_sliding_window"],
+            )
+        )
+
+    prefix_pad_masks = payload["prefix_pad_masks"]
+    if payload.get("prefix_pad_masks_is_bool"):
+        prefix_pad_masks = prefix_pad_masks.to(torch.bool)
+
+    return PrefixContext(
+        past_key_values=VLADensePrefixCache(tuple(kv_layers)),
+        prefix_pad_masks=prefix_pad_masks,
+        prefix_len=int(payload["prefix_len"]),
+        layout=dict(payload["layout"]),
+        cache_key_digest=payload["cache_key_digest"],
+    )

python/sglang/multimodal_gen/runtime/vla/prefix_cache.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections import OrderedDict
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+
+@dataclass
+class PrefixContext:
+    """Request-local observation K/V reused by every action denoise step.
+
+    The optional digest identifies an exact server-level cache entry; suffix K/V
+    is step-dependent and never becomes part of this context.
+    """
+
+    past_key_values: Any
+    prefix_pad_masks: torch.Tensor
+    prefix_len: int
+    layout: dict[str, Any] = field(default_factory=dict)
+    cache_key_digest: str | None = None
+
+
+class VLADensePrefixCache:
+    """a lightweight and naive dense per-layer K/V container for prefix fill and suffix attention.
+
+    Mutable instances collect prefix K/V layer by layer. Read-only instances
+    prepend that fixed K/V to the current suffix K/V without changing storage.
+    """
+
+    def __init__(
+        self,
+        layers: Iterable[tuple[torch.Tensor, torch.Tensor, Any]] | None = None,
+        *,
+        read_only: bool = False,
+    ):
+        # cached_keys, cached_values, sliding_window
+        self.layers = list(layers or ())
+        self.read_only = read_only
+
+    def __iter__(self):
+        return iter(self.layers)
+
+    def __len__(self) -> int:
+        return len(self.layers)
+
+    def __getitem__(self, layer_idx: int):
+        return self.layers[layer_idx]
+
+    def get_seq_length(self) -> int:
+        return 0 if not self.layers else int(self.layers[0][0].shape[-2])
+
+    def get_prefix(self, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor]:
+        prefix_keys, prefix_values, _ = self.layers[layer_idx]
+        return prefix_keys, prefix_values
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """update the cache with fresh kv from each layer, return the appended prefix kv"""
+        if self.read_only:
+            prefix_keys, prefix_values = self.get_prefix(layer_idx)
+            return (
+                torch.cat([prefix_keys, key_states], dim=-2),
+                torch.cat([prefix_values, value_states], dim=-2),
+            )
+
+        if layer_idx == len(self.layers):
+            self.layers.append((key_states, value_states, None))
+            return key_states, value_states
+        if layer_idx > len(self.layers):
+            raise IndexError(f"Invalid VLA prefix cache layer: {layer_idx}")
+        cached_keys, cached_values, sliding_window = self.layers[layer_idx]
+        key_states = torch.cat([cached_keys, key_states], dim=-2)
+        value_states = torch.cat([cached_values, value_states], dim=-2)
+        self.layers[layer_idx] = (key_states, value_states, sliding_window)
+        return key_states, value_states
+
+
+def slice_prefix_context(context: PrefixContext, index: int) -> PrefixContext:
+    return PrefixContext(
+        past_key_values=VLADensePrefixCache(
+            tuple(
+                (
+                    keys[index : index + 1],
+                    values[index : index + 1],
+                    sliding_window,
+                )
+                for keys, values, sliding_window in context.past_key_values
+            )
+        ),
+        prefix_pad_masks=context.prefix_pad_masks[index : index + 1],
+        prefix_len=context.prefix_len,
+        layout=dict(context.layout),
+        cache_key_digest=context.cache_key_digest,
+    )
+
+
+class VLAPrefixCacheManager:
+    """Bounded exact-match LRU for server-level VLA PrefixContext reuse.
+
+    Partial-match prefix cache does not work well VLA scenario (with multiple combinations of keys).
+
+    Request-local denoise reuse does not go through this cache. Partial-prefix
+    K/V reuse is invalid for VLA prefix blocks that use full attention.
+    """
+
+    def __init__(self, max_entries: int = 128):
+        self.max_entries = max(0, int(max_entries))
+        self._cache: OrderedDict[str, PrefixContext] = OrderedDict()
+
+    @staticmethod
+    def make_key(
+        *,
+        model_revision: str,
+        tokenizer_id: str,
+        camera_order: tuple[str, ...],
+        image_hashes: dict[str, str],
+        token_digest: str,
+        token_mask_digest: str,
+        masks: dict[str, bool],
+        positions_version: str,
+        dtype: str,
+        parallel_layout_version: str,
+        cache_namespace: str = "vla",
+    ) -> str:
+        # hash the effective prefix inputs plus runtime compatibility dimensions
+        payload = {
+            "cache_namespace": cache_namespace,
+            "model_revision": model_revision,
+            "tokenizer_id": tokenizer_id,
+            "camera_order": list(camera_order),
+            "image_hashes": image_hashes,
+            "token_digest": token_digest,
+            "token_mask_digest": token_mask_digest,
+            "masks": masks,
+            "positions_version": positions_version,
+            "dtype": dtype,
+            "parallel_layout_version": parallel_layout_version,
+        }
+        serialized = json.dumps(
+            payload,
+            sort_keys=True,
+            separators=(",", ":"),
+            default=str,
+        )
+        return hashlib.sha256(serialized.encode("utf-8")).hexdigest()
+
+    def get(self, key: str) -> PrefixContext | None:
+        context = self._cache.get(key)
+        if context is not None:
+            self._cache.move_to_end(key)
+        return context
+
+    def put(self, key: str, context: PrefixContext) -> None:
+        if self.max_entries == 0:
+            return
+        if len(self._cache) >= self.max_entries and key not in self._cache:
+            self._cache.popitem(last=False)
+        context.cache_key_digest = key
+        self._cache[key] = context
+        self._cache.move_to_end(key)

python/sglang/multimodal_gen/runtime/warmup_request_builder.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Build synthetic diffusion warmup requests.
+"""Build synthetic generation warmup requests.
 
 Default server warmup should cover a representative serving path before the
 first real request, without copying user traffic. It starts from the model's
@@ -17,10 +17,7 @@ from copy import copy
 from typing import Any
 
 from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
-from sglang.multimodal_gen.configs.sample.sampling_params import (
-    DataType,
-    SamplingParams,
-)
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
 from sglang.multimodal_gen.registry import get_pipeline_config_classes
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.server_args import (
@@ -146,7 +143,7 @@ def _fallback_warmup_resolution(server_args: ServerArgs) -> tuple[int, int]:
 
 
 def _is_video_warmup_task(server_args: ServerArgs) -> bool:
-    return server_args.pipeline_config.task_type.data_type() == DataType.VIDEO
+    return server_args.pipeline_config.task_type.is_video_gen()
 
 
 def _warmup_resolution_alignment(server_args: ServerArgs) -> int:
@@ -234,7 +231,7 @@ def _resolve_warmup_num_frames(
     *,
     server_based_warmup: bool,
 ) -> int:
-    num_frames = sampling_defaults.num_frames
+    num_frames = getattr(sampling_defaults, "num_frames", 1)
     if (
         not server_based_warmup
         or not _is_video_warmup_task(server_args)
@@ -247,9 +244,9 @@ def _resolve_warmup_num_frames(
 
 
 def _effective_cfg_scale(sampling_defaults: SamplingParams) -> float | None:
-    if sampling_defaults.true_cfg_scale is not None:
+    if getattr(sampling_defaults, "true_cfg_scale", None) is not None:
         return sampling_defaults.true_cfg_scale
-    return sampling_defaults.guidance_scale
+    return getattr(sampling_defaults, "guidance_scale", None)
 
 
 def _resolve_warmup_steps(
@@ -295,6 +292,8 @@ def should_include_warmup_image(
     server_args: ServerArgs, server_based_warmup: bool
 ) -> bool:
     task_type = server_args.pipeline_config.task_type
+    if not supports_synthetic_warmup(server_args):
+        return False
     if not task_type.accepts_image_input():
         return False
     if task_type.requires_image_input():
@@ -306,6 +305,11 @@ def should_include_warmup_image(
     return True
 
 
+def supports_synthetic_warmup(server_args: ServerArgs) -> bool:
+    task_type = server_args.pipeline_config.task_type
+    return task_type.is_visual_gen() or task_type.is_mesh_gen()
+
+
 def build_warmup_reqs(
     server_args: ServerArgs,
     *,
@@ -315,6 +319,8 @@ def build_warmup_reqs(
     server_based_warmup: bool = False,
 ) -> list[Req]:
     task_type = server_args.pipeline_config.task_type
+    if not supports_synthetic_warmup(server_args):
+        return []
     sampling_defaults = get_model_sampling_defaults(server_args)
 
     if warmup_resolutions is None:
@@ -327,7 +333,7 @@ def build_warmup_reqs(
     else:
         resolutions = [parse_size(resolution) for resolution in warmup_resolutions]
 
-    negative_prompt: Any = sampling_defaults.negative_prompt
+    negative_prompt: Any = getattr(sampling_defaults, "negative_prompt", None)
     cfg_scale = _effective_cfg_scale(sampling_defaults)
     warmup_steps = _resolve_warmup_steps(
         server_args,

python/sglang/multimodal_gen/test/server/gpu_cases.py

@@ -30,6 +30,7 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
     MULTI_FRAME_I2I_sampling_params,
     MULTI_IMAGE_TI2I_sampling_params,
     MULTI_IMAGE_TI2I_UPLOAD_sampling_params,
+    PI05_ACTION_CI_sampling_params,
     SANA_WM_TI2V_CI_sampling_params,
     T2I_sampling_params,
     T2V_sampling_params,
@@ -102,6 +103,16 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
         run_models_api_check=False,
         run_t2v_input_reference_check=False,
     ),
+    DiffusionTestCase(
+        "pi05_action_http",
+        DiffusionServerArgs(
+            model_path="lerobot/pi05_base",
+        ),
+        PI05_ACTION_CI_sampling_params,
+        run_perf_check=False,
+        run_component_accuracy_check=False,
+        run_t2v_input_reference_check=False,
+    ),
     DiffusionTestCase(
         "flux_image_t2i",
         DiffusionServerArgs(model_path=DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST),

python/sglang/multimodal_gen/test/server/test_server_common.py

@@ -7,6 +7,7 @@ Each collected request prints a performance log before validation.
 
 from __future__ import annotations
 
+import json
 import os
 import queue
 import threading
@@ -14,6 +15,7 @@ import time
 from pathlib import Path
 from typing import Any, Callable
 
+import numpy as np
 import openai
 import pytest
 import requests
@@ -47,8 +49,11 @@ from sglang.multimodal_gen.test.test_utils import (
     SGL_TEST_FILES_CI_DATA_REVISION,
     _consistency_gt_filenames,
     _get_consistency_gt_dir,
+    action_gt_exists,
     compare_with_gt,
     extract_key_frames_from_video,
+    get_action_consistency_gt_candidates,
+    get_action_consistency_gt_remote_files,
     get_consistency_gt_candidates,
     get_consistency_gt_remote_files,
     get_consistency_threshold_path,
@@ -56,6 +61,7 @@ from sglang.multimodal_gen.test.test_utils import (
     get_dynamic_server_port,
     gt_exists,
     image_bytes_to_numpy,
+    load_action_consistency_gt,
     load_consistency_gt,
     save_consistency_failure_artifact,
     wait_for_req_perf_record,
@@ -593,6 +599,10 @@ class DiffusionServerBase:
             )
             return
 
+        if case.server_args.modality == "action":
+            self._validate_action_consistency(case, content)
+            return
+
         num_gpus = case.server_args.num_gpus
         is_video = case.server_args.modality == "video"
         output_format = case.sampling_params.output_format
@@ -727,6 +737,88 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION}
             f"max_mean_abs_diff={result.max_mean_abs_diff:.4f})"
         )
 
+    def _extract_action_array(
+        self,
+        payload: dict[str, Any],
+        expected_horizon: int,
+        expected_dim: int,
+    ) -> np.ndarray:
+        action = payload["data"][0]["action"]
+        values = action["values"]
+        assert action["shape"] == [expected_horizon, expected_dim]
+        array = np.asarray(values, dtype=np.float32)
+        assert array.shape == (expected_horizon, expected_dim)
+        assert np.isfinite(array).all()
+        return array
+
+    def _validate_action_consistency(
+        self,
+        case: DiffusionTestCase,
+        content: bytes,
+    ) -> None:
+        payload = json.loads(content.decode("utf-8"))
+        expected_horizon = int(case.sampling_params.extras.get("action_horizon", 50))
+        expected_dim = int(case.sampling_params.extras.get("action_dim", 32))
+        output = self._extract_action_array(payload, expected_horizon, expected_dim)
+
+        num_gpus = case.server_args.num_gpus
+        if not action_gt_exists(case.id, num_gpus):
+            names = ", ".join(get_action_consistency_gt_candidates(case.id, num_gpus))
+            logger.error(f"""
+--- MISSING ACTION GROUND TRUTH DETECTED ---
+GT action JSON not found for '{case.id}'.
+
+Add the expected file to sgl-project/ci-data in diffusion-ci/consistency_gt/sglang_generated/ with naming:
+  Action: {case.id}_{{n}}gpu.json
+
+For this case, expected file(s): {names}
+
+Repository: https://github.com/sgl-project/ci-data (path: diffusion-ci/consistency_gt/sglang_generated/, with optional platform subdirectories such as 5090/)
+Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION}
+""")
+            pytest.fail(
+                f"GT action JSON not found for {case.id}. See logs for instructions to add GT."
+            )
+
+        gt_payload = load_action_consistency_gt(case.id, num_gpus)
+        gt = self._extract_action_array(gt_payload, expected_horizon, expected_dim)
+        abs_diff = np.abs(output - gt)
+        max_abs_diff = float(abs_diff.max())
+        mean_abs_diff = float(abs_diff.mean())
+        max_abs_threshold = float(
+            case.sampling_params.extras.get("action_max_abs_diff_threshold", 0.05)
+        )
+        mean_abs_threshold = float(
+            case.sampling_params.extras.get("action_mean_abs_diff_threshold", 0.005)
+        )
+
+        if max_abs_diff > max_abs_threshold or mean_abs_diff > mean_abs_threshold:
+            gt_remote_info = "\n".join(
+                f"    - {filename}: {url}"
+                for filename, url in get_action_consistency_gt_remote_files(
+                    case.id,
+                    num_gpus,
+                )
+            )
+            pytest.fail(
+                f"Action consistency check failed for {case.id}:\n"
+                f"  max_abs_diff={max_abs_diff:.6f} "
+                f"(threshold {max_abs_threshold:.6f})\n"
+                f"  mean_abs_diff={mean_abs_diff:.6f} "
+                f"(threshold {mean_abs_threshold:.6f})\n"
+                f"  Compared GT files and links:\n{gt_remote_info}"
+            )
+
+        logger.info(
+            "[Consistency] %s: PASSED action GT check "
+            "(shape=%sx%s, max_abs_diff=%.6f, mean_abs_diff=%.6f)",
+            case.id,
+            expected_horizon,
+            expected_dim,
+            max_abs_diff,
+            mean_abs_diff,
+        )
+
     def _save_gt_output(
         self,
         case: DiffusionTestCase,
@@ -749,6 +841,12 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION}
         num_gpus = case.server_args.num_gpus
         is_video = case.server_args.modality == "video"
 
+        if case.server_args.modality == "action":
+            output_path = out_dir / f"{case.id}_{num_gpus}gpu.json"
+            output_path.write_bytes(content)
+            logger.info(f"Saved GT action JSON: {output_path}")
+            return
+
         if is_video:
             # realtime consistency uses websocket raw frames to avoid lossy mp4 drift
             frames = pop_realtime_key_frames(case.id)

python/sglang/multimodal_gen/test/server/test_server_utils.py

@@ -854,6 +854,7 @@ VALIDATOR_REGISTRY = {
     "default": PerformanceValidator,
     "video": VideoPerformanceValidator,
     "mesh": MeshValidator,
+    "action": PerformanceValidator,
 }
 
 
@@ -1501,8 +1502,109 @@ def get_generate_fn(
 
         pytest.fail(f"{case_id}: mesh generation timed out after {max_wait}s")
 
+    def generate_action(case_id, client) -> tuple[str, bytes]:
+        """VLA action generation using /v1/actions/generations."""
+        import numpy as np
+        import requests as http_requests
+
+        extra = dict(sampling_params.extras)
+        action_horizon = int(extra.get("action_horizon", 50))
+        action_dim = int(extra.get("action_dim", 32))
+        state_dim = int(extra.get("state_dim", action_dim))
+        image_size = int(extra.get("image_size", 64))
+        camera_order = tuple(
+            extra.get(
+                "camera_order",
+                ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb"),
+            )
+        )
+
+        def tensor_payload(array):
+            return {
+                "dtype": str(array.dtype),
+                "shape": list(array.shape),
+                "values": array.tolist(),
+            }
+
+        def image_payload(camera_index: int):
+            y = np.arange(image_size, dtype=np.uint16)[:, None]
+            x = np.arange(image_size, dtype=np.uint16)[None, :]
+            image = np.stack(
+                (
+                    (x + camera_index * 17) % 256 + np.zeros_like(y),
+                    (y + camera_index * 29) % 256 + np.zeros_like(x),
+                    (x + y + camera_index * 41) % 256,
+                ),
+                axis=-1,
+            )
+            return tensor_payload(image.astype(np.uint8))
+
+        rng = np.random.default_rng(int(extra.get("seed", 0)))
+        request_id = f"{case_id}-{int(time.time() * 1000)}"
+        payload = {
+            "request_id": request_id,
+            "model": model_path,
+            "input": {
+                "task": sampling_params.prompt or "pick up the blue block",
+                "observation": {
+                    "images": {
+                        camera: image_payload(index)
+                        for index, camera in enumerate(camera_order)
+                    },
+                    "camera_order": list(camera_order),
+                    "state": tensor_payload(
+                        np.linspace(-0.5, 0.5, state_dim, dtype=np.float32)
+                    ),
+                    "noise": tensor_payload(
+                        rng.standard_normal((action_horizon, action_dim)).astype(
+                            np.float32
+                        )
+                    ),
+                },
+            },
+            "parameters": {
+                "action_horizon": action_horizon,
+                "action_dim": action_dim,
+                "num_inference_steps": int(extra.get("num_inference_steps", 2)),
+            },
+            "runtime": {
+                "return_timing": True,
+                "prefix_cache": bool(extra.get("enable_prefix_cache", False)),
+                "cuda_graph": bool(extra.get("enable_cuda_graph", True)),
+                "output_format": "list",
+            },
+        }
+
+        base_url = str(client.base_url).rstrip("/")
+        endpoint = (
+            f"{base_url}/actions/generations"
+            if base_url.endswith("/v1")
+            else f"{base_url}/v1/actions/generations"
+        )
+        response = http_requests.post(endpoint, json=payload, timeout=600)
+        if response.status_code != 200:
+            pytest.fail(f"{case_id}: action generation failed: {response.text}")
+
+        body = response.json()
+        action = body["data"][0]["action"]
+        if action["shape"] != [action_horizon, action_dim]:
+            pytest.fail(
+                f"{case_id}: action shape mismatch: {action['shape']} "
+                f"!= {[action_horizon, action_dim]}"
+            )
+        values = action["values"]
+        if not all(
+            isinstance(value, (int, float)) and np.isfinite(value)
+            for row in values
+            for value in row
+        ):
+            pytest.fail(f"{case_id}: action response contains non-finite values")
+        return body["id"], response.content
+
     if modality == "3d":
         fn = generate_mesh
+    elif modality == "action":
+        fn = generate_action
     elif modality == "video":
         if sampling_params.realtime_num_chunks is not None:
             fn = generate_realtime_video

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -167,7 +167,7 @@ class DiffusionServerArgs:
     """Configuration for a single model/scenario test case."""
 
     model_path: str  # HF repo or local path
-    modality: str | None = None  # auto-inferred: "image" or "video" or "3d"
+    modality: str | None = None  # auto-inferred: "image", "video", "3d", or "action"
 
     custom_validator: str | None = None  # auto-derived unless explicitly overridden
     # resources
@@ -208,6 +208,8 @@ class DiffusionServerArgs:
             self.custom_validator = "video"
         elif self.modality == "3d":
             self.custom_validator = "mesh"
+        elif self.modality == "action":
+            self.custom_validator = "action"
 
 
 @lru_cache(maxsize=None)
@@ -219,6 +221,8 @@ def _infer_modality_from_model_path(model_path: str) -> str:
     task_type = model_info.pipeline_config_cls.task_type
     if task_type == ModelTaskType.I2M:
         return "3d"
+    if task_type.is_action_gen():
+        return "action"
     if task_type.is_image_gen():
         return "image"
     return "video"
@@ -357,6 +361,23 @@ LINGBOT_WORLD_REALTIME_sampling_params = DiffusionSamplingParams(
 )
 
 
+PI05_ACTION_CI_sampling_params = DiffusionSamplingParams(
+    prompt="pick up the blue block",
+    extras={
+        "action_horizon": 50,
+        "action_dim": 32,
+        "state_dim": 32,
+        "image_size": 64,
+        "num_inference_steps": 2,
+        "seed": 0,
+        "enable_prefix_cache": False,
+        "enable_cuda_graph": True,
+        "action_max_abs_diff_threshold": 0.05,
+        "action_mean_abs_diff_threshold": 0.005,
+    },
+)
+
+
 def sample_step_indices(
     step_map: dict[int, float], fractions: Sequence[float]
 ) -> list[int]:
@@ -646,6 +667,8 @@ def get_default_sampling_params_for_model_task(
         return TI2V_sampling_params
     if task_type == ModelTaskType.I2M:
         return HUNYUAN3D_SHAPE_sampling_params
+    if task_type.is_action_gen():
+        return PI05_ACTION_CI_sampling_params
     raise ValueError(f"No default sampling params for model task {task_type!r}")
 
 

python/sglang/multimodal_gen/test/single_test_file/test_pi05_e2e.py

@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import json
+import os
+import statistics
+import time
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+
+pytestmark = pytest.mark.skipif(
+    os.getenv("SGLANG_RUN_PI05_E2E") != "1",
+    reason="set SGLANG_RUN_PI05_E2E=1 to run Pi0.5 GPU e2e tests",
+)
+
+_MODEL_PATH = os.getenv("SGLANG_PI05_E2E_MODEL", "lerobot/pi05_base")
+_CAMERA_ORDER = ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")
+
+
+def _env_int(name: str, default: int) -> int:
+    return int(os.getenv(name, str(default)))
+
+
+def _env_float(name: str) -> float | None:
+    value = os.getenv(name)
+    return None if value is None else float(value)
+
+
+def _image(camera_index: int) -> np.ndarray:
+    height = width = _env_int("SGLANG_PI05_E2E_IMAGE_SIZE", 224)
+    y = np.arange(height, dtype=np.uint16)[:, None]
+    x = np.arange(width, dtype=np.uint16)[None, :]
+    image = np.stack(
+        (
+            (x + camera_index * 17) % 256 + np.zeros_like(y),
+            (y + camera_index * 29) % 256 + np.zeros_like(x),
+            (x + y + camera_index * 41) % 256,
+        ),
+        axis=-1,
+    )
+    return image.astype(np.uint8)
+
+
+def _action_request_kwargs(tag: str) -> dict:
+    action_horizon = _env_int("SGLANG_PI05_E2E_ACTION_HORIZON", 50)
+    action_dim = _env_int("SGLANG_PI05_E2E_ACTION_DIM", 32)
+    rng = np.random.default_rng(_env_int("SGLANG_PI05_E2E_NOISE_SEED", 0))
+    prompt = os.getenv("SGLANG_PI05_E2E_PROMPT", "pick up the blue block")
+    return {
+        "prompt": f"{prompt} [{tag}]",
+        "images": {name: _image(idx) for idx, name in enumerate(_CAMERA_ORDER)},
+        "camera_order": list(_CAMERA_ORDER),
+        "state": np.linspace(
+            -0.5,
+            0.5,
+            _env_int("SGLANG_PI05_E2E_STATE_DIM", 32),
+            dtype=np.float32,
+        ),
+        "noise": rng.standard_normal((action_horizon, action_dim)).astype(np.float32),
+        "action_horizon": action_horizon,
+        "action_dim": action_dim,
+        "num_inference_steps": _env_int("SGLANG_PI05_E2E_NUM_STEPS", 2),
+        "return_timing": True,
+        "enable_prefix_cache": True,
+        "enable_cuda_graph": os.getenv("SGLANG_PI05_E2E_CUDA_GRAPH", "1") != "0",
+    }
+
+
+@pytest.fixture(scope="module")
+def pi05_generator():
+    num_gpus = _env_int("SGLANG_PI05_E2E_NUM_GPUS", 1)
+    kwargs = {
+        "model_path": _MODEL_PATH,
+        "num_gpus": num_gpus,
+        "warmup": False,
+        "trust_remote_code": False,
+    }
+    if num_gpus > 1:
+        kwargs.update(
+            {
+                "sp_degree": _env_int("SGLANG_PI05_E2E_SP_DEGREE", num_gpus),
+                "ulysses_degree": _env_int(
+                    "SGLANG_PI05_E2E_ULYSSES_DEGREE",
+                    num_gpus,
+                ),
+                "ring_degree": _env_int("SGLANG_PI05_E2E_RING_DEGREE", 1),
+            }
+        )
+    generator = DiffGenerator.from_pretrained(local_mode=True, **kwargs)
+    try:
+        yield generator
+    finally:
+        generator.shutdown()
+
+
+def _actions(output: dict) -> np.ndarray:
+    return np.asarray(output["actions"], dtype=np.float32)
+
+
+def _assert_action_output(
+    output: dict, *, expect_cache_hit: bool | None = None
+) -> None:
+    actions = _actions(output)
+    assert actions.shape[0] == _env_int("SGLANG_PI05_E2E_ACTION_HORIZON", 50)
+    expected_output_dim = os.getenv("SGLANG_PI05_E2E_OUTPUT_ACTION_DIM")
+    if expected_output_dim is not None:
+        assert actions.shape[1] == int(expected_output_dim)
+    else:
+        assert 0 < actions.shape[1] <= _env_int("SGLANG_PI05_E2E_ACTION_DIM", 32)
+    assert np.isfinite(actions).all()
+    timings = output.get("timings") or {}
+    assert timings.get("preprocess_ms", 0.0) >= 0.0
+    assert timings.get("prefix_ms", 0.0) >= 0.0
+    assert timings.get("action_denoise_ms", 0.0) > 0.0
+    assert timings.get("postprocess_ms", 0.0) >= 0.0
+
+    cache = output.get("cache") or {}
+    if expect_cache_hit is not None:
+        assert bool(cache.get("hit")) is expect_cache_hit
+
+    parallel = output.get("parallel") or {}
+    num_gpus = _env_int("SGLANG_PI05_E2E_NUM_GPUS", 1)
+    assert bool(parallel.get("split_group", False)) is (num_gpus > 1)
+    if num_gpus > 1:
+        assert int(parallel["world_size"]) == num_gpus
+        assert parallel["prefix_root"] == 0
+        assert parallel["action_root"] == num_gpus - 1
+        assert bool(parallel.get("action_sequence_parallel")) is True
+
+
+def test_pi05_python_action_e2e(pi05_generator):
+    output = pi05_generator.generate_action(_action_request_kwargs("e2e"))
+    _assert_action_output(output)
+
+
+def test_pi05_python_action_consistency(pi05_generator):
+    first = pi05_generator.generate_action(_action_request_kwargs("consistency"))
+    second = pi05_generator.generate_action(_action_request_kwargs("consistency"))
+    _assert_action_output(first, expect_cache_hit=False)
+    _assert_action_output(second, expect_cache_hit=True)
+
+    first_actions = _actions(first)
+    second_actions = _actions(second)
+    np.testing.assert_allclose(
+        first_actions,
+        second_actions,
+        rtol=_env_float("SGLANG_PI05_E2E_CONSISTENCY_RTOL") or 1e-3,
+        atol=_env_float("SGLANG_PI05_E2E_CONSISTENCY_ATOL") or 1e-3,
+    )
+
+    gt_path = os.getenv("SGLANG_PI05_E2E_CONSISTENCY_GT")
+    if gt_path is None:
+        return
+    path = Path(gt_path)
+    if os.getenv("SGLANG_PI05_E2E_UPDATE_GT") == "1":
+        path.parent.mkdir(parents=True, exist_ok=True)
+        np.save(path, first_actions)
+    else:
+        np.testing.assert_allclose(
+            first_actions,
+            np.load(path),
+            rtol=_env_float("SGLANG_PI05_E2E_GT_RTOL") or 1e-2,
+            atol=_env_float("SGLANG_PI05_E2E_GT_ATOL") or 1e-2,
+        )
+
+
+def test_pi05_python_action_perf(pi05_generator):
+    for _ in range(_env_int("SGLANG_PI05_E2E_PERF_WARMUP", 1)):
+        pi05_generator.generate_action(_action_request_kwargs("perf"))
+
+    records = []
+    for _ in range(_env_int("SGLANG_PI05_E2E_PERF_REPEAT", 3)):
+        start = time.perf_counter()
+        output = pi05_generator.generate_action(_action_request_kwargs("perf"))
+        wall_ms = (time.perf_counter() - start) * 1000
+        _assert_action_output(output, expect_cache_hit=True)
+        records.append(
+            {
+                "wall_ms": wall_ms,
+                "timings": output.get("timings") or {},
+                "parallel": output.get("parallel") or {},
+            }
+        )
+
+    wall = [record["wall_ms"] for record in records]
+    denoise = [record["timings"].get("action_denoise_ms", 0.0) for record in records]
+    summary = {
+        "model": _MODEL_PATH,
+        "num_gpus": _env_int("SGLANG_PI05_E2E_NUM_GPUS", 1),
+        "num_steps": _env_int("SGLANG_PI05_E2E_NUM_STEPS", 2),
+        "median_wall_ms": statistics.median(wall),
+        "median_action_denoise_ms": statistics.median(denoise),
+        "records": records,
+    }
+
+    dump_path = os.getenv("SGLANG_PI05_E2E_PERF_DUMP")
+    if dump_path:
+        path = Path(dump_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+
+    max_wall_ms = _env_float("SGLANG_PI05_E2E_MAX_WALL_MS")
+    if max_wall_ms is not None:
+        assert summary["median_wall_ms"] <= max_wall_ms
+    max_denoise_ms = _env_float("SGLANG_PI05_E2E_MAX_ACTION_DENOISE_MS")
+    if max_denoise_ms is not None:
+        assert summary["median_action_denoise_ms"] <= max_denoise_ms

python/sglang/multimodal_gen/test/test_utils.py

@@ -1082,6 +1082,30 @@ def get_consistency_gt_candidates(
     ]
 
 
+def _action_consistency_gt_filenames(case_id: str, num_gpus: int) -> list[str]:
+    case_id = get_consistency_gt_case_id(case_id)
+    return [f"{case_id}_{num_gpus}gpu.json"]
+
+
+def get_action_consistency_gt_candidate_sets(
+    case_id: str,
+    num_gpus: int,
+) -> list[list[str]]:
+    candidates = _action_consistency_gt_filenames(case_id, num_gpus)
+    if _is_ascend_consistency_case(case_id) or current_platform.is_npu():
+        return [candidates]
+    platform = get_consistency_platform()
+    return [[f"{platform}/{candidate}" for candidate in candidates], candidates]
+
+
+def get_action_consistency_gt_candidates(case_id: str, num_gpus: int) -> list[str]:
+    return [
+        candidate
+        for candidate_set in get_action_consistency_gt_candidate_sets(case_id, num_gpus)
+        for candidate in candidate_set
+    ]
+
+
 def get_consistency_gt_remote_files(
     case_id: str, num_gpus: int, is_video: bool, output_format: str | None = None
 ) -> list[tuple[str, str]]:
@@ -1097,6 +1121,19 @@ def get_consistency_gt_remote_files(
     )
 
 
+def get_action_consistency_gt_remote_files(
+    case_id: str, num_gpus: int
+) -> list[tuple[str, str]]:
+    files = _find_remote_action_consistency_gt_files(case_id, num_gpus)
+    if files:
+        return files
+    filenames = get_action_consistency_gt_candidates(case_id, num_gpus)
+    return [
+        (filename, f"{SGL_TEST_FILES_CONSISTENCY_GT_BASE}/{filename}")
+        for filename in filenames
+    ]
+
+
 def _remote_consistency_gt_candidates(
     base_url: str,
     case_id: str,
@@ -1300,6 +1337,35 @@ def _find_remote_consistency_gt_files(
     return []
 
 
+def _find_remote_action_consistency_gt_files(
+    case_id: str,
+    num_gpus: int,
+) -> list[tuple[str, str]]:
+    for filenames in get_action_consistency_gt_candidate_sets(case_id, num_gpus):
+        for base_url in _remote_consistency_gt_base_urls(case_id):
+            candidates = [
+                (filename, f"{base_url}/{filename}") for filename in filenames
+            ]
+            if _is_official_consistency_gt_base_url(base_url):
+                candidates = [
+                    (filename, url)
+                    for filename, url in candidates
+                    if _official_consistency_gt_candidate_is_declared(case_id, filename)
+                ]
+                if not candidates:
+                    continue
+            uncertain_candidate = None
+            for filename, url in candidates:
+                exists = _remote_file_exists(url)
+                if exists is True:
+                    return [(filename, url)]
+                if exists is None and uncertain_candidate is None:
+                    uncertain_candidate = (filename, url)
+            if uncertain_candidate is not None:
+                return [uncertain_candidate]
+    return []
+
+
 def _get_consistency_gt_dir() -> Path | None:
     """Return the local GT directory when configured."""
     d = os.environ.get("SGLANG_CONSISTENCY_GT_DIR")
@@ -1320,6 +1386,13 @@ def _get_consistency_gt_cache_key(
     return f"{platform}:{case_id}:{num_gpus}:{is_video}:{output_format or ''}:{source}"
 
 
+def _get_action_consistency_gt_cache_key(case_id: str, num_gpus: int) -> str:
+    gt_dir = _get_consistency_gt_dir()
+    source = str(gt_dir) if gt_dir is not None else "remote"
+    platform = get_consistency_platform()
+    return f"{platform}:{case_id}:{num_gpus}:action:{source}"
+
+
 def load_consistency_gt(
     case_id: str,
     num_gpus: int,
@@ -1398,6 +1471,60 @@ def load_consistency_gt(
     return loaded_gt
 
 
+def _load_remote_gt_json(url: str) -> dict[str, Any]:
+    last_error: Exception | None = None
+    for _ in range(3):
+        try:
+            resp = requests.get(url, timeout=60)
+            try:
+                if resp.status_code == 200:
+                    return resp.json()
+                last_error = FileNotFoundError(f"GT JSON not found: {url}")
+                if resp.status_code not in (403, 429) and resp.status_code < 500:
+                    break
+            finally:
+                resp.close()
+        except (ValueError, requests.RequestException) as exc:
+            last_error = exc
+    raise FileNotFoundError(f"GT JSON not found: {url}") from last_error
+
+
+def load_action_consistency_gt(case_id: str, num_gpus: int) -> dict[str, Any]:
+    cache_key = _get_action_consistency_gt_cache_key(case_id, num_gpus)
+    cached = _consistency_gt_cache.get(cache_key)
+    if cached is not None:
+        return cached
+
+    gt_dir = _get_consistency_gt_dir()
+    if gt_dir is not None:
+        path = None
+        for fn in get_action_consistency_gt_candidates(case_id, num_gpus):
+            candidate = gt_dir / fn
+            if candidate.exists():
+                path = candidate
+                break
+        if path is None:
+            candidates = get_action_consistency_gt_candidates(case_id, num_gpus)
+            raise FileNotFoundError(
+                f"GT action JSON not found in {gt_dir}. Tried: {', '.join(candidates)}"
+            )
+        with path.open("r", encoding="utf-8") as f:
+            loaded_gt = json.load(f)
+        logger.info("Loaded action GT for %s from %s", case_id, path)
+    else:
+        remote_files = _find_remote_action_consistency_gt_files(case_id, num_gpus)
+        if not remote_files:
+            candidates = get_action_consistency_gt_candidates(case_id, num_gpus)
+            raise FileNotFoundError(
+                f"GT action JSON not found for {case_id}. Tried: {', '.join(candidates)}"
+            )
+        loaded_gt = _load_remote_gt_json(remote_files[0][1])
+        logger.info("Loaded action GT for %s from %s", case_id, remote_files[0][1])
+
+    _consistency_gt_cache[cache_key] = loaded_gt
+    return loaded_gt
+
+
 def load_gt_embeddings(
     case_id: str,
     num_gpus: int,
@@ -1449,6 +1576,26 @@ def gt_exists(
     return found
 
 
+def action_gt_exists(case_id: str, num_gpus: int) -> bool:
+    gt_dir = _get_consistency_gt_dir()
+    if gt_dir is not None:
+        return any(
+            (gt_dir / candidate).exists()
+            for candidate_set in get_action_consistency_gt_candidate_sets(
+                case_id, num_gpus
+            )
+            for candidate in candidate_set
+        )
+
+    cache_key = _get_action_consistency_gt_cache_key(case_id, num_gpus)
+    if cache_key in _gt_exists_remote_cache:
+        return True
+    found = bool(_find_remote_action_consistency_gt_files(case_id, num_gpus))
+    if found:
+        _gt_exists_remote_cache.add(cache_key)
+    return found
+
+
 def extract_key_frames_from_video(
     video_bytes: bytes,
     num_frames: int | None = None,

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

@@ -41,12 +41,17 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import (
 from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import (
     InputValidationStage,
 )
-from sglang.multimodal_gen.runtime.server_warmup import format_warmup_req
+from sglang.multimodal_gen.runtime.server_warmup import (
+    format_warmup_req,
+    should_run_explicit_client_warmup,
+    should_run_synthetic_server_warmup,
+)
 from sglang.multimodal_gen.runtime.warmup_request_builder import (
     DEFAULT_PLACEHOLDER_PROMPT,
     SERVER_WARMUP_IMAGE_FALLBACK_RESOLUTION,
     build_warmup_reqs,
     should_include_warmup_image,
+    supports_synthetic_warmup,
 )
 
 
@@ -642,7 +647,12 @@ class TestWarmupReqCfgParallel(unittest.TestCase):
             ModelTaskType.I2I: True,
             ModelTaskType.I2V: True,
             ModelTaskType.I2M: True,
+            ModelTaskType.VLA_ACTION: False,
         }
+        request_based_expected = {
+            task_type: task_type.accepts_image_input() for task_type in ModelTaskType
+        }
+        request_based_expected[ModelTaskType.VLA_ACTION] = False
 
         for task_type in ModelTaskType:
             server_args = MagicMock()
@@ -655,10 +665,84 @@ class TestWarmupReqCfgParallel(unittest.TestCase):
             )
             self.assertEqual(
                 should_include_warmup_image(server_args, server_based_warmup=False),
-                task_type.accepts_image_input(),
+                request_based_expected[task_type],
                 task_type.name,
             )
 
+    def test_action_pipeline_skips_synthetic_warmup_before_sampling_defaults(self):
+        server_args = MagicMock()
+        server_args.pipeline_config.task_type = ModelTaskType.VLA_ACTION
+
+        with (
+            patch(
+                "sglang.multimodal_gen.runtime.warmup_request_builder.get_model_sampling_defaults"
+            ) as get_defaults,
+            patch(
+                "sglang.multimodal_gen.runtime.warmup_request_builder._resolve_default_warmup_resolution"
+            ) as resolve_resolution,
+        ):
+            reqs = build_warmup_reqs(
+                server_args,
+                warmup_resolutions=None,
+                server_based_warmup=True,
+            )
+
+        self.assertEqual(reqs, [])
+        get_defaults.assert_not_called()
+        resolve_resolution.assert_not_called()
+
+    def test_action_pipeline_disables_synthetic_warmup(self):
+        server_args = MagicMock()
+        server_args.warmup = True
+        server_args.server_warmup = True
+        server_args.warmup_resolutions = ["512x512"]
+        server_args.pipeline_config.task_type = ModelTaskType.VLA_ACTION
+
+        self.assertFalse(supports_synthetic_warmup(server_args))
+        self.assertFalse(should_run_synthetic_server_warmup(server_args))
+        self.assertFalse(should_run_explicit_client_warmup(server_args))
+
+    def test_mesh_pipeline_builds_image_conditioned_warmup(self):
+        server_args = MagicMock()
+        server_args.warmup = True
+        server_args.server_warmup = True
+        server_args.warmup_steps = 1
+        server_args.warmup_resolutions = None
+        server_args.enable_cfg_parallel = False
+        server_args.enable_torch_compile = False
+        server_args.enable_breakable_cuda_graph = False
+        server_args.backend = "native"
+        server_args.pipeline_class_name = None
+        server_args.is_arg_explicitly_set.return_value = False
+        server_args.pipeline_config = SimpleNamespace(
+            task_type=ModelTaskType.I2M,
+            vae_stride=None,
+            vae_scale_factor=None,
+            vae_config=None,
+        )
+
+        with (
+            patch(
+                "sglang.multimodal_gen.runtime.warmup_request_builder.get_model_sampling_defaults",
+                return_value=SamplingParams(width=512, height=512),
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_warmup.is_realtime_serving",
+                return_value=False,
+            ),
+        ):
+            reqs = build_warmup_reqs(
+                server_args,
+                warmup_resolutions=None,
+                warmup_input_path="/tmp/warmup.png",
+                server_based_warmup=True,
+            )
+            self.assertTrue(should_run_synthetic_server_warmup(server_args))
+
+        self.assertEqual(len(reqs), 1)
+        self.assertEqual(reqs[0].data_type, ModelTaskType.I2M.data_type())
+        self.assertEqual(reqs[0].image_path, ["/tmp/warmup.png"])
+
     def test_server_based_warmup_keeps_ti2i_image_input(self):
         server_args = MagicMock()
         server_args.warmup_steps = 1

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

@@ -413,6 +413,37 @@ def test_consistency_gt_case_alias_reuses_canonical_filename(monkeypatch):
     ]
 
 
+def test_action_gt_candidates_prefer_platform_then_default(monkeypatch):
+    monkeypatch.setenv(test_utils.CONSISTENCY_PLATFORM_ENV, "h100")
+
+    assert test_utils.get_action_consistency_gt_candidates("unit_action", 1) == [
+        "h100/unit_action_1gpu.json",
+        "unit_action_1gpu.json",
+    ]
+
+
+def test_remote_action_gt_uses_sglang_generated(monkeypatch):
+    monkeypatch.setenv(test_utils.CONSISTENCY_PLATFORM_ENV, "h100")
+    sglang_prefix = test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE + "/"
+    monkeypatch.setattr(
+        test_utils,
+        "_remote_file_exists",
+        lambda url: url.startswith(sglang_prefix),
+    )
+
+    files = test_utils._find_remote_action_consistency_gt_files("unit_action", 1)
+
+    assert files == [
+        (
+            "h100/unit_action_1gpu.json",
+            (
+                f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}"
+                "/h100/unit_action_1gpu.json"
+            ),
+        )
+    ]
+
+
 def test_threshold_metadata_merges_platform_override():
     metadata = test_utils._merge_threshold_metadata(
         {

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

@@ -0,0 +1,241 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+from types import SimpleNamespace
+
+import numpy as np
+
+from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
+from sglang.multimodal_gen.configs.sample.pi05 import Pi05SamplingParams
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
+from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (
+    action_generation_response,
+    action_metadata,
+    action_raw_response,
+    build_action_sampling_params,
+    pack_msgpack,
+    unpack_msgpack,
+)
+
+
+def _server_args(config: Pi05PipelineConfig | None = None) -> SimpleNamespace:
+    return SimpleNamespace(
+        model_id=None,
+        model_path="lerobot/pi05_base",
+        output_path=None,
+        comfyui_mode=False,
+        num_gpus=1,
+        tp_size=1,
+        sp_degree=1,
+        ulysses_degree=1,
+        ring_degree=1,
+        pipeline_config=config or Pi05PipelineConfig(),
+    )
+
+
+def test_pi05_uses_vla_sampling_params_not_visual_sampling_params():
+    params = Pi05SamplingParams()
+    field_names = {field.name for field in dataclasses.fields(params)}
+
+    assert isinstance(params, VLASamplingParams)
+    assert not isinstance(params, SamplingParams)
+    assert "action_horizon" in field_names
+    assert "action_dim" in field_names
+    assert "height" not in field_names
+    assert "width" not in field_names
+    assert "fps" not in field_names
+    assert "negative_prompt" not in field_names
+    assert "return_frames" not in field_names
+    assert "diffusers_kwargs" not in field_names
+
+
+def test_action_adjust_skips_visual_image_video_logic():
+    params = SamplingParams()
+    params.num_frames = 0
+    params.adjust_frames = True
+    params.return_file_paths_only = True
+
+    params._adjust(_server_args())
+
+    assert params.data_type == DataType.ACTION
+    assert params.num_frames == 1
+    assert params.adjust_frames is False
+    assert params.return_file_paths_only is False
+
+
+def test_action_request_schema_builds_pi05_sampling_params():
+    image = np.zeros((8, 8, 3), dtype=np.uint8)
+    payload = {
+        "request_id": "action-req-1",
+        "model": "lerobot/pi05_base",
+        "input": {
+            "task": "pick up the block",
+            "observation": {
+                "images": {
+                    "base_0_rgb": {
+                        "dtype": "uint8",
+                        "shape": [8, 8, 3],
+                        "values": image.tolist(),
+                    },
+                },
+                "state": {
+                    "dtype": "float32",
+                    "shape": [32],
+                    "values": np.arange(32, dtype=np.float32).tolist(),
+                },
+            },
+        },
+        "parameters": {
+            "action_horizon": 25,
+            "action_dim": 32,
+            "num_inference_steps": 4,
+        },
+        "runtime": {
+            "return_timing": "false",
+            "prefix_cache": False,
+            "cuda_graph": "0",
+        },
+    }
+
+    params = build_action_sampling_params(payload, _server_args())
+
+    assert params.prompt == "pick up the block"
+    assert params.request_id == "action-req-1"
+    assert params.action_horizon == 25
+    assert params.action_dim == 32
+    assert params.num_inference_steps == 4
+    assert not params.return_timing
+    assert not params.enable_prefix_cache
+    assert not params.enable_cuda_graph
+    assert params.return_file_paths_only is False
+    assert params.save_output is False
+    assert set(params.images) == {"base_0_rgb"}
+    assert params.images["base_0_rgb"].shape == (8, 8, 3)
+    assert params.state.shape == (32,)
+
+    vla_state = params.build_request_extra()["vla"]
+    assert vla_state["observation"]["prompt"] == "pick up the block"
+    assert not vla_state["options"]["enable_prefix_cache"]
+
+
+def test_openpi_raw_observation_compatibility_fields_are_normalized():
+    payload = {
+        "task": "push the cube",
+        "observation.images.base_0_rgb": np.ones((4, 4, 3), dtype=np.uint8),
+        "observation.state": {
+            "dtype": "float32",
+            "shape": [32],
+            "data": [0.25] * 32,
+        },
+        "observation.noise": {
+            "dtype": "float32",
+            "shape": [50, 32],
+            "data": np.zeros((50, 32), dtype=np.float32).tolist(),
+        },
+        "enable_pi_prefix_cache": False,
+        "enable_pi_cuda_graph": False,
+    }
+
+    params = build_action_sampling_params(payload, _server_args())
+
+    assert params.prompt == "push the cube"
+    assert set(params.images) == {"base_0_rgb"}
+    assert params.state.shape == (32,)
+    assert params.noise.shape == (50, 32)
+    assert not params.enable_prefix_cache
+    assert not params.enable_cuda_graph
+
+
+def test_action_metadata_reports_policy_shape_and_capabilities():
+    config = Pi05PipelineConfig(
+        image_keys=("front", "wrist"),
+        image_size=(256, 256),
+        state_dim=8,
+        action_horizon=10,
+        action_dim=32,
+        output_action_dim=7,
+        enable_action_cuda_graph=True,
+    )
+
+    metadata = action_metadata(_server_args(config))
+
+    assert metadata["object"] == "action.metadata"
+    assert metadata["policy_family"] == "pi05"
+    assert metadata["input"]["image_keys"] == ["front", "wrist"]
+    assert metadata["input"]["image_size"] == [256, 256]
+    assert metadata["input"]["state_dim"] == 8
+    assert metadata["output"]["action_horizon"] == 10
+    assert metadata["output"]["action_dim"] == 7
+    assert metadata["output"]["padded_action_dim"] == 32
+    assert metadata["runtime"]["materialize_dtype"] == "bf16"
+    assert metadata["runtime"]["enable_autocast"] is True
+    assert metadata["runtime"]["parallelism"]["num_gpus"] == 1
+    assert metadata["runtime"]["parallelism"]["prefix_strategy"] == "tp"
+    assert metadata["runtime"]["parallelism"]["action_strategy"] == "sp"
+    assert metadata["defaults"]["prefix_cache"] is False
+    assert metadata["capabilities"]["realtime_websocket"]
+    assert metadata["capabilities"]["openpi_websocket"]
+
+
+def test_action_generation_response_uses_actual_output_parameters():
+    output = {
+        "request_id": "action-response-1",
+        "actions": [[1.0, 2.0], [3.0, 4.0]],
+        "parameters": {"num_inference_steps": 3},
+        "timings": {"preprocess_ms": 1.5},
+        "cache": {"hit": True},
+        "parallel": {"split_group": False},
+    }
+
+    response = action_generation_response(output, _server_args())
+
+    assert response["id"] == "action-response-1"
+    assert response["object"] == "action.generation"
+    assert response["data"][0]["action"]["shape"] == [2, 2]
+    assert response["data"][0]["action"]["values"] == output["actions"]
+    assert response["usage"]["action_horizon"] == 2
+    assert response["usage"]["action_dim"] == 2
+    assert response["usage"]["denoise_steps"] == 3
+    assert response["usage"]["prefix_cache_hit"] is True
+    assert response["timings"] == output["timings"]
+    assert response["cache"] == output["cache"]
+    assert response["parallel"] == output["parallel"]
+
+
+def test_action_raw_response_preserves_policy_payload_shape():
+    actions = np.arange(6, dtype=np.float32).reshape(2, 3)
+    output = {
+        "actions": actions,
+        "timings": {"preprocess_ms": 1.5},
+    }
+
+    response = action_raw_response(output)
+
+    assert response["actions"] == actions.tolist()
+    assert response["timings"] == output["timings"]
+
+
+def test_action_raw_response_can_preserve_numpy_for_msgpack():
+    actions = np.arange(6, dtype=np.float32).reshape(2, 3)
+
+    response = action_raw_response({"actions": actions}, preserve_numpy=True)
+
+    assert response["actions"] is actions
+
+
+def test_msgpack_roundtrip_preserves_string_keys_and_numpy_payloads():
+    payload = {
+        "task": "pick",
+        "array": np.arange(6, dtype=np.float32).reshape(2, 3),
+        "scalar": np.float32(1.25),
+    }
+
+    decoded = unpack_msgpack(pack_msgpack(payload))
+
+    assert decoded["task"] == "pick"
+    np.testing.assert_array_equal(decoded["array"], payload["array"])
+    assert decoded["scalar"] == payload["scalar"]

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

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from sglang.multimodal_gen.runtime.vla.prefix_cache import (
+    PrefixContext,
+    VLAPrefixCacheManager,
+)
+
+
+def _context() -> PrefixContext:
+    return PrefixContext(
+        past_key_values=("kv",),
+        prefix_pad_masks=torch.ones(1, 3, dtype=torch.bool),
+        prefix_len=3,
+    )
+
+
+def test_pi05_prefix_cache_full_hit_returns_context():
+    manager = VLAPrefixCacheManager(max_entries=4)
+    key = "a"
+    context = _context()
+
+    manager.put(key, context)
+    cached = manager.get(key)
+
+    assert cached is context
+    assert context.cache_key_digest == key
+
+
+def test_pi05_prefix_cache_different_key_misses():
+    manager = VLAPrefixCacheManager(max_entries=4)
+    manager.put("a", _context())
+
+    assert manager.get("b") is None
+
+
+def test_pi05_prefix_cache_zero_capacity_does_not_retain_context():
+    manager = VLAPrefixCacheManager(max_entries=0)
+    context = _context()
+
+    manager.put("a", context)
+
+    assert manager.get("a") is None
+    assert context.cache_key_digest is None
+
+
+def test_pi05_prefix_cache_evicts_least_recently_used_entry():
+    manager = VLAPrefixCacheManager(max_entries=2)
+    context_a = _context()
+    context_b = _context()
+    context_c = _context()
+    manager.put("a", context_a)
+    manager.put("b", context_b)
+    assert manager.get("a") is context_a
+
+    manager.put("c", context_c)
+
+    assert manager.get("a") is context_a
+    assert manager.get("b") is None
+    assert manager.get("c") is context_c

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

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+import sglang.multimodal_gen.runtime.models.vlas.pi05_policy as pi05_policy_module
+from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
+from sglang.multimodal_gen.runtime.models.vlas.pi05_core import (
+    Pi05SiglipAttention,
+    patch_siglip_vision_attention_to_native,
+)
+from sglang.multimodal_gen.runtime.models.vlas.pi05_policy import (
+    Pi05CheckpointManifest,
+    Pi05PolicyModel,
+)
+from sglang.multimodal_gen.runtime.vla.denoise_cuda_graph import (
+    VLADenoiseGraphRunner,
+    _CapturedDenoiseGraph,
+)
+from sglang.multimodal_gen.runtime.vla.parallel import VLASplitGroup
+from sglang.multimodal_gen.runtime.vla.prefix_cache import (
+    PrefixContext,
+    VLADensePrefixCache,
+)
+
+
+def _prefix_context(value: float, digest: str | None) -> PrefixContext:
+    keys = torch.full((1, 1, 2, 4), value)
+    values = torch.full((1, 1, 2, 4), value)
+    return PrefixContext(
+        past_key_values=VLADensePrefixCache(((keys, values, None),)),
+        prefix_pad_masks=torch.ones(1, 2, dtype=torch.bool),
+        prefix_len=2,
+        cache_key_digest=digest,
+    )
+
+
+def test_vla_split_group_marks_all_action_ranks():
+    split = VLASplitGroup(
+        group=SimpleNamespace(world_size=2),
+        prefix_root=0,
+        action_root=1,
+        action_ranks=(0, 1),
+        rank=0,
+    )
+
+    assert split.is_prefix_rank
+    assert split.is_action_rank
+    assert split.uses_action_sp
+
+
+def test_denoise_graph_skips_prefix_copy_for_same_digest(monkeypatch):
+    runner = VLADenoiseGraphRunner(enabled=True)
+    static_context = _prefix_context(1.0, "same")
+    captured = _CapturedDenoiseGraph(
+        graph=object(),
+        static_prefix_context=static_context,
+        static_x_t=torch.empty(1, 2, 4),
+        static_timestep=torch.empty(1),
+        static_output=torch.empty(1, 2, 4),
+        current_context_id=123,
+        current_context_digest="same",
+    )
+
+    def fail_copy(*args, **kwargs):
+        raise AssertionError("PrefixContext should not be copied on digest hit")
+
+    monkeypatch.setattr(
+        "sglang.multimodal_gen.runtime.vla.denoise_cuda_graph._copy_prefix_context_",
+        fail_copy,
+    )
+
+    runner._sync_context_if_needed(captured, _prefix_context(2.0, "same"))
+
+    assert captured.static_prefix_context.past_key_values[0][0].eq(1.0).all()
+
+
+def test_runai_direct_gpu_loader_does_not_reject_split_roles(monkeypatch):
+    class FakeSafeOpen:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, tb):
+            return False
+
+        def keys(self):
+            return ["action.weight"]
+
+    monkeypatch.setattr(
+        pi05_policy_module,
+        "safe_open",
+        lambda *args, **kwargs: FakeSafeOpen(),
+    )
+
+    model = Pi05PolicyModel.__new__(Pi05PolicyModel)
+    model.device = torch.device("cuda")
+    model.runtime_role = "action"
+    model.config = Pi05PipelineConfig()
+    model.manifest = Pi05CheckpointManifest(
+        model_path="fake",
+        safetensor_files=["fake.safetensors"],
+    )
+    model._should_read_source_key = lambda key: True
+    target_state = {
+        "action.weight": SimpleNamespace(device=SimpleNamespace(type="cuda")),
+    }
+
+    assert model._should_stream_weights_to_gpu(target_state, {})
+
+    model.runtime_role = "idle"
+    assert model._should_stream_weights_to_gpu(target_state, {})
+
+
+def test_pi05_loader_maps_unfused_prefix_weights_to_parallel_targets():
+    q_key = (
+        "paligemma_with_expert.paligemma.model.language_model.layers.0."
+        "self_attn.q_proj.weight"
+    )
+    gate_key = (
+        "paligemma_with_expert.paligemma.model.language_model.layers.0."
+        "mlp.gate_proj.weight"
+    )
+
+    assert (
+        "paligemma_with_expert.paligemma.model.language_model.layers.0."
+        "self_attn.qkv_proj.weight",
+        "q",
+    ) in Pi05PolicyModel._candidate_target_weights(q_key)
+    assert (
+        "paligemma_with_expert.paligemma.model.language_model.layers.0."
+        "mlp.gate_up_proj.weight",
+        0,
+    ) in Pi05PolicyModel._candidate_target_weights(gate_key)
+
+
+def test_action_parallel_info_reports_single_rank_without_process_group():
+    model = Pi05PolicyModel.__new__(Pi05PolicyModel)
+    model.runtime_role = "all"
+
+    info = model.action_parallel_info(prefix_context=None)
+
+    assert info == {
+        "split_group": False,
+        "runtime_role": "all",
+        "action_sequence_parallel": False,
+    }
+
+
+class _FakeSiglipAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed_dim = 8
+        self.num_heads = 2
+        self.head_dim = 4
+        self.scale = self.head_dim**-0.5
+        self.dropout = 0.0
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+
+def test_siglip_attention_patch_uses_native_wrapper_once():
+    layer = SimpleNamespace(self_attn=_FakeSiglipAttention())
+    vision_model = SimpleNamespace(encoder=SimpleNamespace(layers=[layer]))
+
+    patch_siglip_vision_attention_to_native(vision_model)
+    first = layer.self_attn
+    patch_siglip_vision_attention_to_native(vision_model)
+
+    assert isinstance(first, Pi05SiglipAttention)
+    assert layer.self_attn is first

python/sglang/utils.py

@@ -32,6 +32,10 @@ from sglang.srt.environ import envs
 logger = logging.getLogger(__name__)
 
 KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS: dict[str, str] = {
+    "lerobot/pi05": "Pi05Pipeline",
+    "lerobot--pi05": "Pi05Pipeline",
+    "pi05": "Pi05Pipeline",
+    "pi0.5": "Pi05Pipeline",
     "hunyuan3d": "Hunyuan3D2Pipeline",
     "flux.2-dev-nvfp4": "Flux2NvfpPipeline",
     "comfy-org/ideogram-4": "Ideogram4Nvfp4Pipeline",