[diffusion] feat: support batching for cosmos3 action generation (#36301)

Signed-off-by: FxxxxU <fu18801374388@163.com>
Signed-off-by: Mick <mickjagger19@icloud.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

docs/cookbook/diffusion/Cosmos/Cosmos3.mdx

@@ -331,6 +331,37 @@ print(action["shape"], action["values"])
 
 Use `GET /v1/actions/metadata` to inspect the action modes, default horizon, padded action dimension, and accepted observation modalities. Msgpack requests and the `/v1/actions/realtime` websocket use the same action envelope.
 
+To batch policy observations inside one request, opt in with a bounded batch size:
+
+```bash Command
+sglang serve \
+  --model-path nvidia/Cosmos3-Nano-Policy-DROID \
+  --num-gpus 1 \
+  --batching-max-size 4
+```
+
+Send one image per observation as a list or `[B, H, W, C]` uint8 array in `input.input_reference`, and either one prompt per image or one scalar prompt to broadcast across the batch. Batched prompts must currently tokenize to the same length because Cosmos3 GEN cross-attention does not mask padded text K/V. All items in one request share the domain, resolution, action horizon, and denoise settings. The standard action envelope returns one `data[i]` item per input, each with action shape `[H, D]`. For a compact msgpack response containing one `[B, H, D]` array, set `runtime.response_format="raw"` and read the top-level `actions` field.
+
+For JSON, `input_reference` can be a list of base64 image payloads. For msgpack, it can be a packed uint8 numpy array directly:
+
+```json JSON
+{
+  "input": {
+    "prompt": ["pick up the block", "close the drawer"],
+    "input_reference": [
+      {"b64_json": "<first-image-base64>"},
+      {"b64_json": "<second-image-base64>"}
+    ]
+  },
+  "parameters": {
+    "action_mode": "policy",
+    "domain_name": "droid_lerobot"
+  }
+}
+```
+
+The batch size cannot exceed `--batching-max-size`; this keeps one request from bypassing the server's configured memory limit. Batching applies to `action_mode="policy"` only. A request seed controls the random stream for the whole batch, so a batched result is deterministic for that request but is not expected to be bit-exact with separately seeded B=1 requests.
+
 `inverse_dynamics` also uses `/v1/actions/generations`; set `action_mode="inverse_dynamics"` and pass an observation video URL or server-local path as `input.observation.video`. Select the embodiment head with `domain_name` or `domain_id`; set `raw_action_dim` explicitly when it cannot be inferred from the domain name.
 
 `forward_dynamics` is intentionally different: it consumes an action array and predicts video, so it remains on `/v1/videos`. Action-producing modes submitted to `/v1/videos` return HTTP 400 with the canonical action endpoint in the error message.

python/sglang/multimodal_gen/runtime/entrypoints/action/cosmos3.py

@@ -0,0 +1,242 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 adapter for the generic action endpoint."""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any
+
+import numpy as np
+from PIL import Image
+
+from sglang.multimodal_gen.configs.sample.cosmos3 import Cosmos3SamplingParams
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+def cosmos3_action_metadata(server_args: ServerArgs) -> dict[str, Any]:
+    pipeline_config = server_args.pipeline_config
+    defaults = Cosmos3SamplingParams()
+    max_batch_size = max(1, int(getattr(server_args, "batching_max_size", 1)))
+    return {
+        "object": "action.metadata",
+        "model": server_args.served_model_name,
+        "model_path": server_args.model_path,
+        "policy_family": "cosmos3",
+        "input": {
+            "modalities": ["image", "video"],
+            "supported_resolutions": [
+                list(resolution) for resolution in defaults.supported_resolutions
+            ],
+            "state_dim": None,
+        },
+        "output": {
+            "action_type": "continuous",
+            "action_horizon": 16,
+            "action_dim": None,
+            "padded_action_dim": pipeline_config.dit_config.arch_config.action_dim,
+            "dtype": "float32",
+        },
+        "runtime": {
+            "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,
+            }
+        },
+        "defaults": {
+            "action_mode": "policy",
+            "action_horizon": 16,
+            "num_inference_steps": defaults.num_inference_steps,
+            "height": 480,
+            "width": 832,
+            "fps": 5,
+        },
+        "capabilities": {
+            "action_modes": ["policy", "inverse_dynamics"],
+            "realtime_websocket": True,
+            "openpi_websocket": False,
+            "batch_inputs": max_batch_size > 1,
+            "max_batch_size": max_batch_size,
+            "batched_action_modes": ["policy"],
+            "multiple_candidates": False,
+        },
+    }
+
+
+def _images_from_observation(observation: dict[str, Any]) -> list[Any]:
+    image = None
+    for name in ("image", "image_path", "input_reference"):
+        if name in observation:
+            image = observation[name]
+            break
+
+    if image is None:
+        images = observation.get("images")
+        if images is None or (isinstance(images, dict) and not images):
+            return []
+        if not isinstance(images, dict) or len(images) != 1:
+            raise ValueError(
+                "Cosmos3 action input accepts one image field; use a list or "
+                "a [B, H, W, C] array in that field for batched observations"
+            )
+        image = next(iter(images.values()))
+
+    if isinstance(image, (list, tuple)):
+        images = list(image)
+    elif isinstance(image, np.ndarray) and image.ndim == 4:
+        images = list(image)
+    else:
+        images = [image]
+
+    normalized_images: list[Any] = []
+    for item in images:
+        if not isinstance(item, np.ndarray):
+            normalized_images.append(item)
+            continue
+        if item.dtype != np.uint8:
+            raise ValueError("Cosmos3 observation image arrays must use uint8 dtype")
+        if item.ndim not in (2, 3):
+            raise ValueError(
+                "Cosmos3 observation image arrays must have shape [H, W] "
+                f"or [H, W, C], got {tuple(item.shape)}"
+            )
+        normalized_images.append(Image.fromarray(item))
+    return normalized_images
+
+
+def _action_prompt(prompt: Any, batch_size: int) -> str | list[str]:
+    if isinstance(prompt, str):
+        return prompt if batch_size == 1 else [prompt] * batch_size
+    if not isinstance(prompt, (list, tuple)) or not prompt:
+        raise ValueError("Cosmos3 action prompt must be a string or non-empty list")
+    if not all(isinstance(item, str) for item in prompt):
+        raise ValueError("Cosmos3 action prompt list must contain only strings")
+    prompts = list(prompt)
+    if len(prompts) == 1 and batch_size > 1:
+        prompts *= batch_size
+    if len(prompts) != batch_size:
+        raise ValueError(
+            "Cosmos3 batched action input requires one prompt per image, got "
+            f"{len(prompts)} prompt(s) and {batch_size} image(s)"
+        )
+    return prompts[0] if batch_size == 1 else prompts
+
+
+def build_cosmos3_action_sampling_params(
+    payload: dict[str, Any],
+    observation: dict[str, Any],
+    server_args: ServerArgs,
+    sampling_params_cls: type[Cosmos3SamplingParams],
+) -> Cosmos3SamplingParams:
+    parameters = dict(payload.get("parameters") or {})
+    options = {**observation, **parameters}
+    action_mode = str(options.get("action_mode", "policy")).strip().lower()
+    if action_mode == "forward_dynamics":
+        raise ValueError(
+            "Cosmos3 forward_dynamics produces video; use /v1/videos instead"
+        )
+    if action_mode not in ("policy", "inverse_dynamics"):
+        raise ValueError(
+            "Cosmos3 action endpoint supports action_mode='policy' or "
+            "'inverse_dynamics'"
+        )
+
+    action_horizon = options.get("action_horizon")
+    num_frames = options.get("num_frames")
+    if action_horizon is None and num_frames is None:
+        action_horizon = 16
+    if action_horizon is not None:
+        action_horizon = int(action_horizon)
+        if action_horizon <= 0:
+            raise ValueError("action_horizon must be a positive integer")
+        expected_num_frames = action_horizon + 1
+        if num_frames is not None and int(num_frames) != expected_num_frames:
+            raise ValueError(
+                "Cosmos3 requires num_frames == action_horizon + 1, got "
+                f"num_frames={num_frames}, action_horizon={action_horizon}"
+            )
+        num_frames = expected_num_frames
+    else:
+        num_frames = int(num_frames)
+        if num_frames <= 1:
+            raise ValueError("Cosmos3 action num_frames must be greater than 1")
+    if (num_frames - 1) % 4 != 0:
+        raise ValueError(
+            "Cosmos3 action_horizon must be divisible by 4 so num_frames "
+            "is compatible with the temporal VAE"
+        )
+
+    images = _images_from_observation(observation)
+    video_path = options.get("video_path") or observation.get("video")
+    if action_mode == "policy" and not images:
+        raise ValueError("Cosmos3 policy input requires an observation image")
+    if action_mode == "inverse_dynamics" and video_path is None:
+        raise ValueError("Cosmos3 inverse_dynamics input requires an observation video")
+    if images and video_path is not None:
+        raise ValueError("Cosmos3 action requests accept either an image or a video")
+    batch_size = len(images) if images else 1
+    max_batch_size = max(1, int(getattr(server_args, "batching_max_size", 1)))
+    if batch_size > max_batch_size:
+        raise ValueError(
+            f"Cosmos3 action batch size {batch_size} exceeds "
+            f"--batching-max-size={max_batch_size}"
+        )
+    image_path = None if not images else images[0] if batch_size == 1 else images
+
+    domain_id = options.get("domain_id")
+    domain_name = options.get("domain_name")
+    raw_action_dim = options.get("raw_action_dim")
+    if domain_id is None and not domain_name:
+        raise ValueError("Cosmos3 action requests require domain_name or domain_id")
+    if domain_id is not None and not domain_name and raw_action_dim is None:
+        raise ValueError("raw_action_dim is required when only domain_id is provided")
+
+    prompt = observation.get("prompt")
+    if prompt is None:
+        prompt = observation.get("task", "")
+    if action_mode != "policy" and not isinstance(prompt, str):
+        raise ValueError("Cosmos3 inverse_dynamics prompt must be a string")
+    prompt = _action_prompt(prompt, batch_size)
+    sampling_kwargs = {
+        "request_id": payload.get("request_id") or payload.get("id"),
+        "prompt": prompt,
+        "image_path": image_path,
+        "video_path": video_path,
+        "action_mode": action_mode,
+        "domain_id": domain_id,
+        "domain_name": domain_name,
+        "raw_action_dim": raw_action_dim,
+        "action_fps": options.get("action_fps"),
+        "action_view_point": options.get("action_view_point", "ego_view"),
+        "action_normalization": options.get("action_normalization", "quantile"),
+        "action_stats_path": server_args.pipeline_config.action_stats_path,
+        "num_frames": num_frames,
+        "fps": int(options.get("fps", 5)),
+        "height": int(options.get("height", 480)),
+        "width": int(options.get("width", 832)),
+        "num_inference_steps": int(options.get("num_inference_steps", 35)),
+        "guidance_scale": float(options.get("guidance_scale", 1.0)),
+        "seed": int(options.get("seed", 42)),
+        "flow_shift": options.get("flow_shift"),
+        "max_sequence_length": options.get("max_sequence_length"),
+        "condition_frame_indexes": options.get("condition_frame_indexes"),
+        "condition_video_keep": options.get("condition_video_keep", "first"),
+        "use_duration_template": False,
+        "use_system_prompt": False,
+        "use_guardrails": options.get("use_guardrails"),
+        "save_output": False,
+        "return_file_paths_only": False,
+        "return_frames": False,
+    }
+    supported_fields = {field.name for field in dataclasses.fields(sampling_params_cls)}
+    sampling_params = sampling_params_cls(
+        **{
+            name: value
+            for name, value in sampling_kwargs.items()
+            if name in supported_fields and value is not None
+        }
+    )
+    sampling_params._adjust(server_args)
+    return sampling_params

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

@@ -17,6 +17,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config
 from sglang.multimodal_gen.configs.sample.action import ActionSamplingParams
 from sglang.multimodal_gen.configs.sample.cosmos3 import Cosmos3SamplingParams
 from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.action.cosmos3 import (
+    build_cosmos3_action_sampling_params,
+    cosmos3_action_metadata,
+)
 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
@@ -112,7 +116,12 @@ def _normalize_observation(observation: dict[str, Any]) -> dict[str, Any]:
         }
     for name in ("image", "image_path", "input_reference"):
         if name in normalized:
-            normalized[name] = _normalize_image_value(normalized[name])
+            value = normalized[name]
+            normalized[name] = (
+                [_normalize_image_value(item) for item in value]
+                if isinstance(value, (list, tuple))
+                else _normalize_image_value(value)
+            )
     state = normalized.get("state")
     if isinstance(state, dict):
         normalized["state"] = _decode_tensor_payload(state)
@@ -148,51 +157,7 @@ def images_from_observation(
 def action_metadata(server_args: ServerArgs) -> dict[str, Any]:
     pipeline_config = server_args.pipeline_config
     if isinstance(pipeline_config, Cosmos3Config):
-        defaults = Cosmos3SamplingParams()
+        return cosmos3_action_metadata(server_args)
-        return {
-            "object": "action.metadata",
-            "model": server_args.served_model_name,
-            "model_path": server_args.model_path,
-            "policy_family": "cosmos3",
-            "input": {
-                "modalities": ["image", "video"],
-                "supported_resolutions": [
-                    list(resolution) for resolution in defaults.supported_resolutions
-                ],
-                "state_dim": None,
-            },
-            "output": {
-                "action_type": "continuous",
-                "action_horizon": 16,
-                "action_dim": None,
-                "padded_action_dim": pipeline_config.dit_config.arch_config.action_dim,
-                "dtype": "float32",
-            },
-            "runtime": {
-                "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,
-                }
-            },
-            "defaults": {
-                "action_mode": "policy",
-                "action_horizon": 16,
-                "num_inference_steps": defaults.num_inference_steps,
-                "height": 480,
-                "width": 832,
-                "fps": 5,
-            },
-            "capabilities": {
-                "action_modes": ["policy", "inverse_dynamics"],
-                "realtime_websocket": True,
-                "openpi_websocket": False,
-                "batch_inputs": False,
-                "multiple_candidates": False,
-            },
-        }
 
     policy_family = getattr(
         pipeline_config,
@@ -423,134 +388,6 @@ def _build_action_model_sampling_params(
     return sp
 
 
-def _cosmos3_image_from_observation(observation: dict[str, Any]) -> Any:
-    image = None
-    for name in ("image", "image_path", "input_reference"):
-        if name in observation:
-            image = observation[name]
-            break
-
-    if image is None:
-        images = observation.get("images")
-        if not images:
-            return None
-        if not isinstance(images, dict) or len(images) != 1:
-            raise ValueError(
-                "Cosmos3 policy input requires exactly one observation image"
-            )
-        image = next(iter(images.values()))
-    if isinstance(image, np.ndarray):
-        if image.dtype != np.uint8:
-            raise ValueError("Cosmos3 observation image arrays must use uint8 dtype")
-        return Image.fromarray(image)
-    return image
-
-
-def _build_cosmos3_action_sampling_params(
-    payload: dict[str, Any],
-    server_args: ServerArgs,
-    sampling_params_cls: type[Cosmos3SamplingParams],
-) -> Cosmos3SamplingParams:
-    observation = _action_request_to_observation(payload)
-    parameters = dict(payload.get("parameters") or {})
-    options = {**observation, **parameters}
-    action_mode = str(options.get("action_mode", "policy")).strip().lower()
-    if action_mode == "forward_dynamics":
-        raise ValueError(
-            "Cosmos3 forward_dynamics produces video; use /v1/videos instead"
-        )
-    if action_mode not in ("policy", "inverse_dynamics"):
-        raise ValueError(
-            "Cosmos3 action endpoint supports action_mode='policy' or "
-            "'inverse_dynamics'"
-        )
-
-    action_horizon = options.get("action_horizon")
-    num_frames = options.get("num_frames")
-    if action_horizon is None and num_frames is None:
-        action_horizon = 16
-    if action_horizon is not None:
-        action_horizon = int(action_horizon)
-        if action_horizon <= 0:
-            raise ValueError("action_horizon must be a positive integer")
-        expected_num_frames = action_horizon + 1
-        if num_frames is not None and int(num_frames) != expected_num_frames:
-            raise ValueError(
-                "Cosmos3 requires num_frames == action_horizon + 1, got "
-                f"num_frames={num_frames}, action_horizon={action_horizon}"
-            )
-        num_frames = expected_num_frames
-    else:
-        num_frames = int(num_frames)
-        if num_frames <= 1:
-            raise ValueError("Cosmos3 action num_frames must be greater than 1")
-    if (num_frames - 1) % 4 != 0:
-        raise ValueError(
-            "Cosmos3 action_horizon must be divisible by 4 so num_frames "
-            "is compatible with the temporal VAE"
-        )
-
-    image_path = _cosmos3_image_from_observation(observation)
-    video_path = options.get("video_path") or observation.get("video")
-    if action_mode == "policy" and image_path is None:
-        raise ValueError("Cosmos3 policy input requires an observation image")
-    if action_mode == "inverse_dynamics" and video_path is None:
-        raise ValueError("Cosmos3 inverse_dynamics input requires an observation video")
-    if image_path is not None and video_path is not None:
-        raise ValueError("Cosmos3 action requests accept either an image or a video")
-
-    domain_id = options.get("domain_id")
-    domain_name = options.get("domain_name")
-    raw_action_dim = options.get("raw_action_dim")
-    if domain_id is None and not domain_name:
-        raise ValueError("Cosmos3 action requests require domain_name or domain_id")
-    if domain_id is not None and not domain_name and raw_action_dim is None:
-        raise ValueError("raw_action_dim is required when only domain_id is provided")
-
-    prompt = observation.get("prompt") or observation.get("task") or ""
-    sampling_kwargs = {
-        "request_id": payload.get("request_id") or payload.get("id"),
-        "prompt": prompt,
-        "image_path": image_path,
-        "video_path": video_path,
-        "action_mode": action_mode,
-        "domain_id": domain_id,
-        "domain_name": domain_name,
-        "raw_action_dim": raw_action_dim,
-        "action_fps": options.get("action_fps"),
-        "action_view_point": options.get("action_view_point", "ego_view"),
-        "action_normalization": options.get("action_normalization", "quantile"),
-        "action_stats_path": server_args.pipeline_config.action_stats_path,
-        "num_frames": num_frames,
-        "fps": int(options.get("fps", 5)),
-        "height": int(options.get("height", 480)),
-        "width": int(options.get("width", 832)),
-        "num_inference_steps": int(options.get("num_inference_steps", 35)),
-        "guidance_scale": float(options.get("guidance_scale", 1.0)),
-        "seed": int(options.get("seed", 42)),
-        "flow_shift": options.get("flow_shift"),
-        "max_sequence_length": options.get("max_sequence_length"),
-        "condition_frame_indexes": options.get("condition_frame_indexes"),
-        "condition_video_keep": options.get("condition_video_keep", "first"),
-        "use_duration_template": False,
-        "use_system_prompt": False,
-        "use_guardrails": options.get("use_guardrails"),
-        "save_output": False,
-        "return_file_paths_only": False,
-        "return_frames": False,
-    }
-    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 and value is not None
-        }
-    )
-    sp._adjust(server_args)
-    return sp
-
-
 def build_action_sampling_params(
     payload: dict[str, Any],
     server_args: ServerArgs,
@@ -563,8 +400,9 @@ def build_action_sampling_params(
             sampling_params_cls,
         )
     if issubclass(sampling_params_cls, Cosmos3SamplingParams):
-        return _build_cosmos3_action_sampling_params(
+        return build_cosmos3_action_sampling_params(
             payload,
+            _action_request_to_observation(payload),
             server_args,
             sampling_params_cls,
         )
@@ -594,23 +432,44 @@ def action_generation_response(
     preserve_numpy: bool = False,
 ) -> dict[str, Any]:
     actions = output["actions"]
-    if isinstance(actions, np.ndarray):
+    action_array = np.asarray(actions)
-        action_shape = list(actions.shape)
+    if any(size == 0 for size in action_array.shape):
-        action_values = actions if preserve_numpy else actions.tolist()
+        raise ValueError(
-    else:
+            "action output dimensions must be non-zero, got "
-        horizon = len(actions) if isinstance(actions, list) else 0
+            f"{tuple(action_array.shape)}"
-        action_dim = len(actions[0]) if horizon and isinstance(actions[0], list) else 0
+        )
-        action_shape = [horizon, action_dim]
+    if action_array.ndim == 2:
-        action_values = actions
+        action_array = action_array[None]
-    action = {
+    elif action_array.ndim != 3:
-        "type": "continuous",
+        raise ValueError(
-        "dtype": "float32",
+            "action output must have shape [H, D] or [B, H, D], got "
-        "shape": action_shape,
+            f"{tuple(action_array.shape)}"
-        "values": action_values,
+        )
-    }
+
-    for name in ("action_mode", "domain_id", "raw_action_dim"):
+    data = []
-        if output.get(name) is not None:
+    for input_index, action_values in enumerate(action_array):
-            action[name] = output[name]
+        action_shape = list(action_values.shape)
+        if not preserve_numpy:
+            action_values = action_values.tolist()
+        action = {
+            "type": "continuous",
+            "dtype": "float32",
+            "shape": action_shape,
+            "values": action_values,
+        }
+        for name in ("action_mode", "domain_id", "raw_action_dim"):
+            if output.get(name) is not None:
+                action[name] = output[name]
+        data.append(
+            {
+                "index": input_index,
+                "input_index": input_index,
+                "candidate_index": 0,
+                "action": action,
+            }
+        )
+
+    action_shape = data[0]["action"]["shape"]
 
     pipeline_config = server_args.pipeline_config
     if isinstance(pipeline_config, Cosmos3Config):
@@ -623,15 +482,9 @@ def action_generation_response(
         "object": "action.generation",
         "created": int(time.time()),
         "model": server_args.served_model_name,
-        "data": [
+        "data": data,
-            {
-                "index": 0,
-                "input_index": 0,
-                "candidate_index": 0,
-                "action": action,
-            }
-        ],
         "usage": {
+            "batch_size": len(data),
             "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(

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

@@ -899,8 +899,17 @@ def prepare_request(
     if diffusers_kwargs and "max_sequence_length" in diffusers_kwargs:
         req.max_sequence_length = diffusers_kwargs["max_sequence_length"]
 
-    if not isinstance(req.prompt, str):
+    action_prompt = (
-        raise TypeError(f"`prompt` must be a string, but got {type(req.prompt)}")
+        req.data_type == DataType.ACTION
+        and isinstance(req.prompt, list)
+        and bool(req.prompt)
+        and all(isinstance(item, str) for item in req.prompt)
+    )
+    if not isinstance(req.prompt, str) and not action_prompt:
+        raise TypeError(
+            "`prompt` must be a string, or a non-empty list of strings for "
+            f"batched action requests, but got {type(req.prompt)}"
+        )
 
     req_width = getattr(req, "width", None)
     req_height = getattr(req, "height", None)

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

@@ -41,6 +41,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_action import (
     ACTION_MODE_FORWARD_DYNAMICS,
     ACTION_MODE_INVERSE_DYNAMICS,
+    ACTION_MODE_POLICY,
     ACTION_MODES,
     EMBODIMENT_TO_DOMAIN_ID,
     build_action_prompt,
@@ -141,7 +142,9 @@ def _pil_to_normalized_tensor(image: PIL.Image.Image) -> torch.Tensor:
 class Cosmos3ImagePreprocessStage(PipelineStage):
     """Load, aspect-resize, and center-crop the conditioning input.
 
-    For I2V: writes ``[1, 3, H, W]`` to ``batch.preprocessed_image``.
+    For I2V: writes ``[1, 3, H, W]`` to ``batch.preprocessed_image``. Batched
+    policy requests write ``[B, 3, H, W]``; regular visual generation remains
+    single-image conditioned.
     For V2V: writes ``[1, 3, T_in, H, W]`` to ``batch.preprocessed_video``.
     No-op for T2V / T2I.
     """
@@ -153,9 +156,14 @@ class Cosmos3ImagePreprocessStage(PipelineStage):
 
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         image_path = batch.image_path
-        if isinstance(image_path, list):
-            image_path = image_path[0] if image_path else None
         video_path = batch.video_path
+        is_action_policy = (
+            batch.data_type == DataType.ACTION
+            and getattr(batch.sampling_params, "action_mode", None)
+            == ACTION_MODE_POLICY
+        )
+        if isinstance(image_path, list) and not is_action_policy:
+            image_path = image_path[0] if image_path else None
         if isinstance(video_path, list):
             video_path = video_path[0] if video_path else None
 
@@ -168,10 +176,23 @@ class Cosmos3ImagePreprocessStage(PipelineStage):
         target_h, target_w = batch.height, batch.width
 
         if image_path is not None:
-            image = load_image(image_path)
+            image_sources = (
-            image = _resize_crop_pil(image, target_w, target_h)
+                list(image_path)
-            batch.preprocessed_image = _pil_to_normalized_tensor(image).unsqueeze(0)
+                if isinstance(image_path, (list, tuple))
-            self.log_info(f"Preprocessed conditioning image to {target_w}x{target_h}")
+                else [image_path]
+            )
+            if not image_sources:
+                raise ValueError("Cosmos3 I2V image list is empty")
+            tensors: list[torch.Tensor] = []
+            for src in image_sources:
+                image = load_image(src)
+                image = _resize_crop_pil(image, target_w, target_h)
+                tensors.append(_pil_to_normalized_tensor(image))
+            batch.preprocessed_image = torch.stack(tensors, dim=0).contiguous()
+            self.log_info(
+                f"Preprocessed {len(tensors)} conditioning image(s) to "
+                f"{target_w}x{target_h}"
+            )
             return batch
 
         if isinstance(video_path, str) and video_path:
@@ -265,7 +286,7 @@ class Cosmos3TokenizationStage(PipelineStage):
 
     def _tokenize_prompt(
         self,
-        text: str,
+        text: str | list[str],
         max_sequence_length: int,
         device: torch.device,
         use_system_prompt: bool = False,
@@ -273,58 +294,72 @@ class Cosmos3TokenizationStage(PipelineStage):
     ) -> tuple[torch.Tensor, torch.Tensor, int]:
         """Tokenize a prompt using Qwen2 chat template.
 
-        Returns (input_ids, attention_mask, seq_len) as [1, S] tensors.
+        Returns (input_ids, attention_mask, seq_len) as [B, S] tensors.
         """
-        conversations = []
+        texts = text if isinstance(text, (list, tuple)) else [text]
-        if use_system_prompt:
+        if not texts:
-            conversations.append(
+            raise ValueError("Cosmos3 prompt batch must not be empty")
-                {
+        input_id_lists: list[list[int]] = []
-                    "role": "system",
+        attention_mask_lists: list[list[int]] = []
-                    "content": system_prompt or COSMOS3_VIDEO_SYSTEM_PROMPT,
+        seq_lens: list[int] = []
-                }
-            )
-        conversations.append({"role": "user", "content": text})
-
-        result = self.tokenizer.apply_chat_template(
-            conversations,
-            tokenize=True,
-            add_generation_prompt=True,
-        )
-
-        # Handle different return types from apply_chat_template
-        # Fast tokenizer returns BatchEncoding, slow tokenizer returns list[int]
-        if hasattr(result, "input_ids"):
-            # BatchEncoding from fast tokenizer
-            token_ids = list(result.input_ids)
-        elif isinstance(result, list):
-            # Already a list from slow tokenizer
-            token_ids = list(result)
-        else:
-            raise TypeError(
-                f"Unexpected return type from apply_chat_template: {type(result)}"
-            )
-
-        # Reserve room for the two special tokens (EOS + vision_start) so the
-        # final length cannot exceed ``max_sequence_length``.
-        token_ids = token_ids[: max_sequence_length - 2]
-
-        # Add EOS and vision_start tokens
-        token_ids.append(self.tokenizer.eos_token_id)
-        vision_start_id = self.tokenizer.convert_tokens_to_ids("<|vision_start|>")
-        if vision_start_id is not None:
-            token_ids.append(vision_start_id)
-
-        seq_len = len(token_ids)
-
-        # Pad to max_sequence_length
-        pad_len = max_sequence_length - seq_len
-        attention_mask = [1] * seq_len + [0] * pad_len
         pad_token_id = self.tokenizer.pad_token_id or 0
-        token_ids = token_ids + [pad_token_id] * pad_len
+        vision_start_id = self.tokenizer.convert_tokens_to_ids("<|vision_start|>")
+        for text_item in texts:
+            conversations = []
+            if use_system_prompt:
+                conversations.append(
+                    {
+                        "role": "system",
+                        "content": system_prompt or COSMOS3_VIDEO_SYSTEM_PROMPT,
+                    }
+                )
+            conversations.append({"role": "user", "content": text_item})
 
-        input_ids = torch.tensor([token_ids], dtype=torch.long, device=device)
+            result = self.tokenizer.apply_chat_template(
-        attention_mask = torch.tensor([attention_mask], dtype=torch.long, device=device)
+                conversations,
-        return input_ids, attention_mask, seq_len
+                tokenize=True,
+                add_generation_prompt=True,
+            )
+            # Handle different return types from apply_chat_template
+            # Fast tokenizer returns BatchEncoding, slow tokenizer returns list[int]
+            if hasattr(result, "input_ids"):
+                # BatchEncoding from fast tokenizer
+                token_ids = list(result.input_ids)
+            elif isinstance(result, list):
+                # Already a list from slow tokenizer
+                token_ids = list(result)
+            else:
+                raise TypeError(
+                    f"Unexpected return type from apply_chat_template: {type(result)}"
+                )
+
+            # Reserve room for the two special tokens (EOS + vision_start) so the
+            # final length cannot exceed ``max_sequence_length``.
+            token_ids = token_ids[: max_sequence_length - 2]
+            # Add EOS and vision_start tokens
+            token_ids.append(self.tokenizer.eos_token_id)
+            if vision_start_id is not None:
+                token_ids.append(vision_start_id)
+
+            seq_len = len(token_ids)
+            pad_len = max_sequence_length - seq_len
+            attention_mask = [1] * seq_len + [0] * pad_len
+            token_ids = token_ids + [pad_token_id] * pad_len
+            input_id_lists.append(token_ids)
+            attention_mask_lists.append(attention_mask)
+            seq_lens.append(seq_len)
+
+        if len(set(seq_lens)) != 1:
+            raise ValueError(
+                "Cosmos3 batched prompts must tokenize to the same length because "
+                "GEN cross-attention does not mask padded text K/V; split prompts "
+                f"into equal-length batches instead (lengths={seq_lens})"
+            )
+        input_ids = torch.tensor(input_id_lists, dtype=torch.long, device=device)
+        attention_mask = torch.tensor(
+            attention_mask_lists, dtype=torch.long, device=device
+        )
+        return input_ids, attention_mask, seq_lens[0]
 
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         """Tokenize prompt and negative prompt."""
@@ -383,6 +418,9 @@ class Cosmos3TokenizationStage(PipelineStage):
             self.log_info(f"Prompt with duration: '{prompt}'")
 
         # Tokenize prompts
+        if isinstance(prompt, list) and not isinstance(negative_prompt, list):
+            negative_prompt = [negative_prompt] * len(prompt)
+
         cond_ids, cond_mask, cond_seq_len = self._tokenize_prompt(
             prompt, max_sequence_length, device, use_system_prompt, system_prompt
         )
@@ -473,8 +511,13 @@ class Cosmos3LatentPreparationStage(PipelineStage):
         height_latent = batch.height // vae_scale_factor_spatial
         width_latent = batch.width // vae_scale_factor_spatial
 
+        if batch.preprocessed_image is not None:
+            batch_dim = int(batch.preprocessed_image.shape[0])
+        else:
+            batch_dim = 1
+
         shape = (
-            1,
+            batch_dim,
             num_channels_latents,
             num_latent_frames,
             height_latent,
@@ -523,7 +566,7 @@ class Cosmos3LatentPreparationStage(PipelineStage):
 
             condition_latents = torch.zeros_like(noise)
             condition_mask = torch.zeros(
-                1, 1, num_latent_frames, 1, 1, device=device, dtype=dtype
+                batch_dim, 1, num_latent_frames, 1, 1, device=device, dtype=dtype
             )
             for idx in cond_indexes:
                 src = min(idx, cond_latent.shape[2] - 1)
@@ -637,6 +680,11 @@ class Cosmos3LatentPreparationStage(PipelineStage):
         action_offset = 1 if action_chunk_size == num_frames - 1 else 0
 
         domain_id = self._resolve_domain_id(batch)
+        batch_dim = (
+            int(batch.raw_latent_shape[0])
+            if getattr(batch, "raw_latent_shape", None)
+            else 1
+        )
         raw_action_dim = getattr(sp, "raw_action_dim", None)
         if raw_action_dim is None:
             embodiment = getattr(sp, "domain_name", None)
@@ -678,7 +726,7 @@ class Cosmos3LatentPreparationStage(PipelineStage):
             if raw_action_dim is None:
                 raise ValueError(f"action_mode={mode!r} requires --raw-action-dim.")
             clean_action = torch.zeros(
-                1, action_chunk_size, action_dim, device=device, dtype=dtype
+                batch_dim, action_chunk_size, action_dim, device=device, dtype=dtype
             )
 
         raw_action_dim = int(raw_action_dim)
@@ -690,13 +738,13 @@ class Cosmos3LatentPreparationStage(PipelineStage):
         # condition_mask marks clean (given) action tokens. forward_dynamics
         # conditions on the whole action sequence; the others denoise it fully.
         condition_mask = torch.zeros(
-            1, action_chunk_size, 1, device=device, dtype=dtype
+            batch_dim, action_chunk_size, 1, device=device, dtype=dtype
         )
         if mode == ACTION_MODE_FORWARD_DYNAMICS:
             condition_mask[:] = 1.0
 
         noise = torch.randn(
-            1,
+            batch_dim,
             action_chunk_size,
             action_dim,
             generator=generator,
@@ -709,7 +757,7 @@ class Cosmos3LatentPreparationStage(PipelineStage):
 
         batch.action_latents = action_latents
         batch.extra["action_domain_ids"] = torch.tensor(
-            [domain_id], dtype=torch.long, device=device
+            [domain_id] * batch_dim, dtype=torch.long, device=device
         )
         batch.extra["action_velocity_mask"] = 1.0 - condition_mask
         batch.extra["action_condition_latents"] = clean_action
@@ -1092,7 +1140,8 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         )
 
         for i, t in progress_bar:
-            timestep = t.unsqueeze(0) if t.dim() == 0 else t
+            batch_dim = batch.latents.shape[0] if batch.latents is not None else 1
+            timestep = t.unsqueeze(0).expand(batch_dim) if t.dim() == 0 else t
             # Outside the CFG window the effective scale collapses to 1.0,
             # which reduces CFG to the cond branch (cfg-parallel safe).
             effective_scale = (
@@ -1399,7 +1448,7 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         latents_batched = torch.cat([latents, latents], dim=0)
         text_ids_batched = torch.cat([uncond_text_ids, cond_text_ids], dim=0)
         text_mask_batched = torch.cat([uncond_text_mask, cond_text_mask], dim=0)
-        timestep_batched = timestep.expand(2)
+        timestep_batched = torch.cat([timestep, timestep], dim=0)
         mask_batched = (
             torch.cat([noisy_frame_mask, noisy_frame_mask], dim=0)
             if noisy_frame_mask is not None
@@ -1645,9 +1694,12 @@ class Cosmos3DecodingStage(PipelineStage):
         if batch.data_type == DataType.ACTION:
             if action_pred is None:
                 raise RuntimeError("Cosmos3 action request produced no action tensor")
+            payload_actions = (
+                action_pred[0] if action_pred.shape[0] == 1 else action_pred
+            )
             payload = {
                 "request_id": batch.request_id,
-                "actions": action_pred[0].numpy(),
+                "actions": payload_actions.numpy(),
                 "action_mode": action_metadata["action_mode"],
                 "domain_id": action_metadata["action_domain_id"],
                 "raw_action_dim": action_metadata["action_raw_action_dim"],

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

@@ -123,31 +123,42 @@ def canonical_aspect_ratio(width: int, height: int) -> str:
 
 
 def build_action_prompt(
-    description: str,
+    description: str | list[str],
     view_point: str,
     num_frames: int,
     fps: float,
     height: int,
     width: int,
-) -> str:
+) -> str | list[str]:
     """Render the structured JSON action caption the action checkpoints expect."""
     duration_seconds = num_frames / fps
     minutes, secs = divmod(round(duration_seconds), 60)
-    if description and description[-1] not in ".!?":
+    if isinstance(description, (list, tuple)):
-        description = description + "."
+        descriptions = [str(d) for d in description]
-    prompt = {
+    else:
-        "cinematography": {
+        descriptions = [description]
-            "framing": VIEWPOINT_TEMPLATES.get(
+
-                view_point, VIEWPOINT_TEMPLATES["ego_view"]
+    prompts = []
-            )
+    for desc in descriptions:
-        },
+        if desc and desc[-1] not in ".!?":
-        "actions": [{"time": f"0:00-{minutes}:{secs:02d}", "description": description}],
+            desc = desc + "."
-        "duration": f"{int(duration_seconds)}s",
+        prompt = {
-        "fps": float(fps),
+            "cinematography": {
-        "resolution": {"H": int(height), "W": int(width)},
+                "framing": VIEWPOINT_TEMPLATES.get(
-        "aspect_ratio": canonical_aspect_ratio(int(width), int(height)),
+                    view_point, VIEWPOINT_TEMPLATES["ego_view"]
-    }
+                )
-    return json.dumps(prompt)
+            },
+            "actions": [{"time": f"0:00-{minutes}:{secs:02d}", "description": desc}],
+            "duration": f"{int(duration_seconds)}s",
+            "fps": float(fps),
+            "resolution": {"H": int(height), "W": int(width)},
+            "aspect_ratio": canonical_aspect_ratio(int(width), int(height)),
+        }
+        prompts.append(json.dumps(prompt))
+
+    if isinstance(description, (list, tuple)):
+        return prompts
+    return prompts[0]
 
 
 def load_action_stats(

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

@@ -8,6 +8,7 @@ import unittest
 from unittest import mock
 
 import torch
+from PIL import Image
 
 from sglang.multimodal_gen.configs.models.dits.cosmos3video import (
     _build_cosmos3_param_names_mapping,
@@ -51,6 +52,7 @@ from sglang.multimodal_gen.runtime.models.dits.cosmos3video import (
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import (
     Cosmos3DecodingStage,
+    Cosmos3DenoisingStage,
     Cosmos3ImagePreprocessStage,
     Cosmos3LatentPreparationStage,
     Cosmos3TimestepPreparationStage,
@@ -70,7 +72,7 @@ def _apply(mapping_fn, key):
     return mapping_fn(key)
 
 
-def _cosmos3_server_args(config=None):
+def _cosmos3_server_args(config=None, batching_max_size=1):
     return types.SimpleNamespace(
         model_id=None,
         model_path="nvidia/Cosmos3-Nano",
@@ -84,6 +86,7 @@ def _cosmos3_server_args(config=None):
         sp_degree=1,
         ulysses_degree=1,
         ring_degree=1,
+        batching_max_size=batching_max_size,
         pipeline_config=config or Cosmos3Config(),
     )
 
@@ -497,6 +500,31 @@ class TestCosmos3SamplingParamsDataType(unittest.TestCase):
 
 
 class TestCosmos3ActionEndpoint(unittest.TestCase):
+    @staticmethod
+    def _policy_payload(
+        batch_size=2,
+        prompt="pick up the block",
+        *,
+        tensor_payload=False,
+    ):
+        images = torch.zeros(batch_size, 8, 8, 3, dtype=torch.uint8).numpy()
+        input_reference = (
+            {
+                "dtype": "uint8",
+                "shape": list(images.shape),
+                "values": images.tolist(),
+            }
+            if tensor_payload
+            else images
+        )
+        return {
+            "input": {"prompt": prompt, "input_reference": input_reference},
+            "parameters": {
+                "action_mode": "policy",
+                "domain_name": "droid_lerobot",
+            },
+        }
+
     def test_policy_request_builds_action_sampling_params(self):
         image = torch.zeros(8, 8, 3, dtype=torch.uint8).numpy()
         payload = {
@@ -535,6 +563,49 @@ class TestCosmos3ActionEndpoint(unittest.TestCase):
         self.assertEqual(params.seed, 7)
         self.assertEqual(params.image_path.size, (8, 8))
 
+    def test_batched_policy_request_preserves_input_pairing(self):
+        payload = self._policy_payload(
+            prompt=["pick up the block", "close the drawer"],
+            tensor_payload=True,
+        )
+
+        params = build_action_sampling_params(
+            payload, _cosmos3_server_args(batching_max_size=2)
+        )
+
+        self.assertEqual(params.prompt, ["pick up the block", "close the drawer"])
+        self.assertEqual(len(params.image_path), 2)
+        self.assertTrue(
+            all(isinstance(image, Image.Image) for image in params.image_path)
+        )
+
+    def test_batched_policy_request_broadcasts_scalar_prompt(self):
+        params = build_action_sampling_params(
+            self._policy_payload(), _cosmos3_server_args(batching_max_size=2)
+        )
+
+        self.assertEqual(params.prompt, ["pick up the block"] * 2)
+
+    def test_batched_policy_request_rejects_cardinality_mismatch(self):
+        with self.assertRaisesRegex(ValueError, "one prompt per image"):
+            build_action_sampling_params(
+                self._policy_payload(prompt=["one", "two", "three"]),
+                _cosmos3_server_args(batching_max_size=3),
+            )
+
+    def test_batched_policy_request_honors_server_batch_limit(self):
+        with self.assertRaisesRegex(ValueError, "--batching-max-size=1"):
+            build_action_sampling_params(self._policy_payload(), _cosmos3_server_args())
+
+    def test_single_item_image_batch_keeps_single_input_contract(self):
+        params = build_action_sampling_params(
+            self._policy_payload(batch_size=1, prompt="pick"),
+            _cosmos3_server_args(),
+        )
+
+        self.assertEqual(params.prompt, "pick")
+        self.assertIsInstance(params.image_path, Image.Image)
+
     def test_inverse_dynamics_maps_video_input(self):
         payload = {
             "input": {
@@ -554,6 +625,21 @@ class TestCosmos3ActionEndpoint(unittest.TestCase):
         self.assertEqual(params.video_path, "observation.mp4")
         self.assertEqual(params.num_frames, 61)
 
+    def test_inverse_dynamics_rejects_prompt_batch(self):
+        payload = {
+            "input": {
+                "prompt": ["first", "second"],
+                "video": "observation.mp4",
+            },
+            "parameters": {
+                "action_mode": "inverse_dynamics",
+                "domain_name": "droid_lerobot",
+            },
+        }
+
+        with self.assertRaisesRegex(ValueError, "prompt must be a string"):
+            build_action_sampling_params(payload, _cosmos3_server_args())
+
     def test_forward_dynamics_is_rejected_by_action_endpoint(self):
         payload = {
             "input": {
@@ -567,7 +653,7 @@ class TestCosmos3ActionEndpoint(unittest.TestCase):
             build_action_sampling_params(payload, _cosmos3_server_args())
 
     def test_metadata_describes_cosmos_action_contract(self):
-        metadata = action_metadata(_cosmos3_server_args())
+        metadata = action_metadata(_cosmos3_server_args(batching_max_size=4))
 
         self.assertEqual(metadata["model"], "cosmos3-production")
         self.assertEqual(metadata["policy_family"], "cosmos3")
@@ -575,6 +661,15 @@ class TestCosmos3ActionEndpoint(unittest.TestCase):
         self.assertEqual(metadata["output"]["action_horizon"], 16)
         self.assertEqual(metadata["output"]["padded_action_dim"], 64)
         self.assertFalse(metadata["capabilities"]["openpi_websocket"])
+        self.assertTrue(metadata["capabilities"]["batch_inputs"])
+        self.assertEqual(metadata["capabilities"]["max_batch_size"], 4)
+        self.assertEqual(metadata["capabilities"]["batched_action_modes"], ["policy"])
+
+    def test_metadata_keeps_batching_opt_in(self):
+        metadata = action_metadata(_cosmos3_server_args())
+
+        self.assertFalse(metadata["capabilities"]["batch_inputs"])
+        self.assertEqual(metadata["capabilities"]["max_batch_size"], 1)
 
     def test_action_response_includes_cosmos_metadata(self):
         output = {
@@ -595,6 +690,33 @@ class TestCosmos3ActionEndpoint(unittest.TestCase):
         self.assertEqual(action["raw_action_dim"], 10)
         self.assertEqual(response["usage"]["denoise_steps"], 30)
 
+    def test_batched_action_response_emits_one_data_item_per_input(self):
+        output = {
+            "request_id": "cosmos-action-batch",
+            "actions": torch.arange(24, dtype=torch.float32).reshape(2, 4, 3).numpy(),
+            "action_mode": "policy",
+            "domain_id": 8,
+            "raw_action_dim": 3,
+        }
+
+        response = action_generation_response(output, _cosmos3_server_args())
+
+        self.assertEqual(len(response["data"]), 2)
+        self.assertEqual(response["data"][0]["action"]["shape"], [4, 3])
+        self.assertEqual(response["data"][1]["input_index"], 1)
+        self.assertEqual(response["usage"]["batch_size"], 2)
+        self.assertEqual(response["usage"]["action_horizon"], 4)
+        self.assertEqual(response["usage"]["action_dim"], 3)
+
+    def test_action_response_rejects_empty_batch(self):
+        output = {
+            "request_id": "cosmos-action-empty",
+            "actions": torch.empty(0, 4, 3).numpy(),
+        }
+
+        with self.assertRaisesRegex(ValueError, "dimensions must be non-zero"):
+            action_generation_response(output, _cosmos3_server_args())
+
     def test_action_decode_skips_vae(self):
         class FailIfDecoded:
             def decode(self, _latents):
@@ -629,6 +751,35 @@ class TestCosmos3ActionEndpoint(unittest.TestCase):
         self.assertEqual(output.output[0]["domain_id"], 8)
         self.assertEqual(output.action_pred.shape, (1, 4, 3))
 
+    def test_batched_action_decode_keeps_batch_dimension(self):
+        stage = Cosmos3DecodingStage.__new__(Cosmos3DecodingStage)
+        stage.vae = mock.Mock()
+        stage.sound_tokenizer = None
+        stage._guardrails = False
+        stage.log_info = lambda *_args, **_kwargs: None
+        batch = types.SimpleNamespace(
+            data_type=DataType.ACTION,
+            action_latents=torch.arange(48, dtype=torch.float32).reshape(2, 4, 6),
+            extra={
+                "raw_action_dim": 3,
+                "action_domain_ids": torch.tensor([8, 8]),
+            },
+            sampling_params=Cosmos3SamplingParams(
+                prompt=["one", "two"],
+                action_mode="policy",
+                domain_name="droid_lerobot",
+            ),
+            request_id="cosmos-action-batch",
+            num_inference_steps=30,
+            num_frames=5,
+            metrics=None,
+        )
+
+        output = stage.forward(batch, types.SimpleNamespace(vae_cpu_offload=False))
+
+        self.assertEqual(output.output[0]["actions"].shape, (2, 4, 3))
+        self.assertEqual(output.action_pred.shape, (2, 4, 3))
+
 
 class TestCosmos3ModelResolution(unittest.TestCase):
     """Verify Cosmos3 checkpoints resolve to the native SGLang pipeline."""
@@ -972,10 +1123,13 @@ class TestCosmos3ActionLatentPrep(unittest.TestCase):
         cls.device = torch.device("cpu")
         cls.dtype = torch.float32
 
-    def _run(self, num_frames=17, **sp_kwargs):
+    def _run(self, num_frames=17, batch_size=1, **sp_kwargs):
         sp = Cosmos3SamplingParams(prompt="t", num_frames=num_frames, **sp_kwargs)
         batch = types.SimpleNamespace(
-            sampling_params=sp, num_frames=num_frames, extra={}
+            sampling_params=sp,
+            num_frames=num_frames,
+            raw_latent_shape=(batch_size, 48, 1, 1, 1),
+            extra={},
         )
         gen = torch.Generator(device=self.device).manual_seed(0)
         self.stage._prepare_action_latents(batch, gen, self.device, self.dtype)
@@ -1010,6 +1164,17 @@ class TestCosmos3ActionLatentPrep(unittest.TestCase):
         # padding dims beyond raw_action_dim start at zero.
         self.assertTrue(torch.all(batch.action_latents[:, :, 10:] == 0))
 
+    def test_batched_policy_prepares_one_action_stream_per_observation(self):
+        batch = self._run(
+            batch_size=3,
+            action_mode="policy",
+            domain_name="droid_lerobot",
+        )
+
+        self.assertEqual(tuple(batch.action_latents.shape), (3, 16, 64))
+        self.assertEqual(tuple(batch.extra["action_domain_ids"].shape), (3,))
+        self.assertEqual(tuple(batch.extra["action_velocity_mask"].shape), (3, 16, 1))
+
     def test_inverse_dynamics_denoises_from_noise(self):
         batch = self._run(
             num_frames=61,
@@ -1041,6 +1206,113 @@ class TestCosmos3ActionLatentPrep(unittest.TestCase):
             self._run(action_mode="teleport", domain_id=0)
 
 
+class TestCosmos3BatchedActionStages(unittest.TestCase):
+    class _Tokenizer:
+        pad_token_id = 0
+        eos_token_id = 1
+
+        @staticmethod
+        def convert_tokens_to_ids(_token):
+            return 2
+
+        @staticmethod
+        def apply_chat_template(conversations, **_kwargs):
+            text = conversations[-1]["content"]
+            return list(range(3, 3 + len(text.split())))
+
+    @staticmethod
+    def _preprocess_images(data_type, action_mode=None):
+        stage = Cosmos3ImagePreprocessStage.__new__(Cosmos3ImagePreprocessStage)
+        stage.log_info = lambda *_args, **_kwargs: None
+        seen = []
+        batch = types.SimpleNamespace(
+            image_path=["first.png", "second.png"],
+            video_path=None,
+            height=16,
+            width=16,
+            data_type=data_type,
+            sampling_params=types.SimpleNamespace(action_mode=action_mode),
+            preprocessed_image=None,
+        )
+
+        def fake_load_image(path):
+            seen.append(path)
+            return Image.new("RGB", (16, 16))
+
+        with mock.patch(
+            "sglang.multimodal_gen.runtime.pipelines_core.stages."
+            "model_specific_stages.cosmos3.load_image",
+            side_effect=fake_load_image,
+        ):
+            stage.forward(batch, types.SimpleNamespace())
+        return seen, batch.preprocessed_image
+
+    def test_tokenization_rejects_unequal_prompt_lengths(self):
+        stage = Cosmos3TokenizationStage.__new__(Cosmos3TokenizationStage)
+        stage.tokenizer = self._Tokenizer()
+
+        with self.assertRaisesRegex(ValueError, "same length"):
+            stage._tokenize_prompt(
+                ["pick block", "close the top drawer"],
+                max_sequence_length=16,
+                device=torch.device("cpu"),
+            )
+
+    def test_tokenization_preserves_equal_length_prompt_batch(self):
+        stage = Cosmos3TokenizationStage.__new__(Cosmos3TokenizationStage)
+        stage.tokenizer = self._Tokenizer()
+
+        input_ids, attention_mask, seq_len = stage._tokenize_prompt(
+            ["pick block", "push cube"],
+            max_sequence_length=16,
+            device=torch.device("cpu"),
+        )
+
+        self.assertEqual(tuple(input_ids.shape), (2, 16))
+        self.assertEqual(tuple(attention_mask.shape), (2, 16))
+        self.assertEqual(seq_len, 4)
+
+    def test_cfg_duplicates_batched_timesteps(self):
+        stage = Cosmos3DenoisingStage.__new__(Cosmos3DenoisingStage)
+        captured = {}
+
+        def fake_run_transformer(**kwargs):
+            captured.update(kwargs)
+            return kwargs["latents"]
+
+        stage._run_transformer = fake_run_transformer
+        latents = torch.zeros(3, 1, 1, 1, 1)
+        text_ids = torch.ones(3, 4, dtype=torch.long)
+        text_mask = torch.ones_like(text_ids)
+
+        output = stage._predict_noise_cfg_batched(
+            latents=latents,
+            timestep=torch.ones(3),
+            cond_text_ids=text_ids,
+            cond_text_mask=text_mask,
+            uncond_text_ids=text_ids,
+            uncond_text_mask=text_mask,
+            video_shape=(1, 1, 1),
+            fps=20.0,
+            guidance_scale=2.0,
+        )
+
+        self.assertEqual(tuple(captured["timestep"].shape), (6,))
+        self.assertEqual(tuple(output.shape), tuple(latents.shape))
+
+    def test_visual_i2v_keeps_single_conditioning_image(self):
+        seen, image = self._preprocess_images(DataType.VIDEO)
+
+        self.assertEqual(seen, ["first.png"])
+        self.assertEqual(tuple(image.shape), (1, 3, 16, 16))
+
+    def test_action_preprocess_stacks_all_images(self):
+        seen, image = self._preprocess_images(DataType.ACTION, "policy")
+
+        self.assertEqual(seen, ["first.png", "second.png"])
+        self.assertEqual(tuple(image.shape), (2, 3, 16, 16))
+
+
 class TestCosmos3ModalitySamplingParams(unittest.TestCase):
     """Sound / V2V / action sampling-param fields and defaults."""