Cosmos3-edge PR 6: Cosmos3-Edge as a cosmos3 config variant, with streaming rollout and action loop

[merceod]

Adds nvidia/Cosmos3-Edge (4B) to the cosmos3 family as a config-driven variant, registry cosmos3_edge and cosmos3_edge_droid. No fork. The dense Nemotron backbone (relu2 MLP, Nemotron RMSNorm ordering, no text QK-norm, k_norm_und_for_gen), the 480p defaults, the native flow schedule and the diffusers-0.40 i2v conditioning recipe are Cosmos3Config knobs read from the checkpoint. The Edge reasoner (understanding tower + SigLIP2-style vision encoder) is served as a VLM on /v1/chat/completions from the same transformer instance and KV pool as the DiT, with its decode step compiled before the CUDA-graph capture. Also registers the 4-step distilled Super checkpoints (cosmos3_super_t2i_4step, cosmos3_super_i2v_4step).

New for serving a world model is the streaming rollout. A long clip is generated window by window (8 latent = 29 pixel frames at 480p), each finished window leaves the denoise loop on a streaming edge to a decoder node in its own partition, and the client receives frames while the next window is still denoising. chained mode re-pins the previous window's tail as clean frames, kv mode commits the window's K/V into the paged cache so later windows attend to it block-causally, with the context bounded by the retention policy from #293. A session_id keeps that world state across requests. The action policy runs over the /generate/ws WebSocket from #294.

Depends on #292 #293 #294 #295 #296. Their commits are carried on this branch until they merge (a rebase drops them), so merge order is those five first, then this one. Parity against diffusers 0.40 and transformers 5.17 (fixed-seed PSNR per frame, engine-cache harness, reasoner greedy tokens) lives in the tests under mstar/model/cosmos3/tests/.

Benchmarks

H100 80GB, one GPU of one node, every system back-to-back on that GPU, warmup excluded, n=3 medians (2026-09-23/24; clocks, driver and versions in the run's environment.txt). 832x480 unless noted, 20 UniPC steps, guidance 6, seed 0. M* is configs/cosmos3_edge.yaml (the DROID row configs/cosmos3_edge_droid.yaml) at 9a5dafb9.

System (version)	streaming rollout, 241 frames i2v: time to first frame / frames per s	i2v / t2v 121 frames, wall s	t2i 640x640, s	DROID action loop, 32 actions per call (4 steps, guidance 3, shift 5)	reasoner TTFT p50 / decode tok/s per stream at 1 / 4 / 8 / 32 streams (image prompt, 256 tokens)
M*	kv 1.87 s / 14.0, chained 1.82 s / 14.3 (whole clip 18.6 s)	7.50 / 7.29	0.44	454 ms per chunk = 70 actions/s at 832x480 (113 with two in flight), 362 ms = 87 at 640x360	50 ms / 435, 75 / 290 (1033 aggregate), 164 / 209 (1383), 1381 / 98 (1829)
vLLM-Omni 0.28.0	no streaming, whole clip 19.3 s	8.06 / 7.99	not exposed for this model	1.37 s per chunk = 23 actions/s (0.90 s at 640x360)	n/a (generator recipe)
vLLM 0.29 (reasoner)	n/a	n/a	n/a	n/a	18 ms / 434, 30 / 425 (1621), 42 / 414 (3097), 127 / 364 (9534)
SGLang 0.5.19 (SGLang-Diffusion)	no streaming, whole clip not run	8.04 / 8.03 (server-side 7.18 / 7.13)	0.89	0.67 s per chunk = 48 actions/s	n/a
diffusers 0.40	n/a	11.3 (latent loop only)	n/a	reference only (used for parity)	n/a
TensorRT-LLM 1.2.1	no Cosmos3 in VisualGen 1.2.1

The reasoner is the one row M* loses at more than one stream. The GPU is not the reason: the decode graph runs at the weight-streaming floor, and a profile shows the step at 32 streams spending ~8 of its 10 ms in per-row CPU work in the worker (postprocess, output routing, one device-to-host copy per row) plus an eager 25 ms vision prefill that stalls every stream. Both are engine-side follow-ups (the per-row work as its own PR), not properties of the model integration.

Also measured earlier (2026-09-17/18, two GPUs of one node at once, n=2, not redone under the one-GPU protocol): Nano at 832x480x121 with 35 steps 30.6 / 30.4 s on M* vs 30.9 / 30.8 on vLLM-Omni (i2v / t2v); the Nano 241-frame rollout streams its first window at 7.0 s where vLLM-Omni returns the clip whole after 76.4 s; Super 4-step distilled on TP2 does t2i 1024x1024 in 0.45 s and i2v 832x480x121 in 6.5 s.

Reproduce with benchmark/cosmos3/reproduce_edge.sh (set SNAP, MSTAR, VLLM_OMNI_PY, VLLM_PY, SGLANG first):

reproduce_edge.sh serve-mstar 0 8100            # M* server (cosmos3_edge.yaml)
reproduce_edge.sh serve-mstar-droid 0 8101      # M* DROID policy server
reproduce_edge.sh serve-vllm-omni 0 8200        # vLLM-Omni generator baseline
reproduce_edge.sh serve-vllm-reasoner 0 8300    # vLLM reasoner baseline
reproduce_edge.sh serve-sglang 0 8400           # SGLang-Diffusion baseline
reproduce_edge.sh bench-generator 8100 8200 <image.jpg>
reproduce_edge.sh bench-sglang 8400 <image.jpg>
reproduce_edge.sh bench-stream 8100 <image.jpg>
reproduce_edge.sh bench-action 8101 8200 8400 <image.jpg>
reproduce_edge.sh bench-reasoner 8100 8300 <image.png>