minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: FxxxxU <fu18801374388@163.com>
Signed-off-by: Mick <mickjagger19@icloud.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
XuFu
2026-08-26 22:15:03 +08:00
committed by GitHub
co-authored by Mick
parent dfc40e0efe
commit 924aeee59c
7 changed files with 759 additions and 289 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/cookbook/diffusion/Cosmos/Cosmos3.mdx
+31
View File
@@ -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.