[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 (declared by the window's commit step). A session_id keeps that world state across requests: the DiT and the decoder keep it in a SessionStore that pins a session while a request runs (never evicted, a second request on it is refused), ages idle ones out after session_timeout_s, and drops one on end_session. The action policy runs over the /generate/ws WebSocket from #294.

Depends on #293 #294 #295 #296 (#292 merged). Their commits, review fixes included, are carried on this branch until they merge (a rebase drops them), so merge order is those four 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-27; 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 716bd83c.

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 / 16 / 32 streams (image prompt, 256 tokens)
M*	kv 1.89 s / 13.8, chained 1.81 s / 14.2 (whole clip 18.6 s)	7.49 / 7.26	0.49	467 ms per chunk = 69 actions/s at 832x480 (112 with two in flight), 379 ms = 84 at 640x360	54 ms / 437, 118 / 302 (1034 aggregate), 282 / 240 (1398), 365 / 154 (1860), 1542 / 105 (1827)
vLLM-Omni 0.28.0	no streaming, whole clip 19.2 s	8.04 / 7.95	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	19 ms / 433, 29 / 424 (1623), 42 / 412 (3085), 66 / 400 (5762), 109 / 361 (9724)
SGLang 0.5.19 (SGLang-Diffusion)	no streaming, whole clip not run	8.04 / 8.03 (server-side 7.23 / 7.14)	0.91	0.69 s per chunk = 46 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, spread over the routing stages (route, register and send outputs, collect outputs, prepare inputs, ~30 us per row each) plus an eager 25 ms vision prefill that stalls every stream. Batching the stop check's host copies (#334, closed) changed nothing measurable there and cost 1-2 ms per step on small image tiers, so it was reverted after this run (the table's tree carried it; the video and action rows do not touch that path, and the t2i row can only have improved). The fix for the reasoner tail is a batched loop-back path in the worker, an engine design item, not a property of the model integration. The t2i number is the eager 640x640 tier, 20 launch-bound steps: it read 0.44 s on an idle node (2026-09-23) and 0.49 s on a node shared with seven other jobs; the GPU-bound rows did not move.

Two servers side by side on the same GPU, 2026-09-27: session-store stress (a second request on a live session is refused, live sessions are never evicted, sessions time out and can be ended), 869-frame kv and chained rollouts at 14 fps with flat memory, cancel storms, a mixed t2i+t2v+i2v+chat+action round, and reasoner cancel rounds all clean. Parity against diffusers 0.40 in fp32: i2v 56.8 dB (frame 0 exact), t2v 39.0 dB.

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>

[NSagan271]

I have several concerns about the handling of persistent sessions, which make the case for giving M* primitives for handling sessions instead of purely putting it into model code:

• The decoder and DiT also don't agree on the LRU ordering: the DiT promotes on read at submodules.py:701, the decoder doesn't at :2737
• If there are 8 in-flight requests with live sessions and a 9th is added, then one of the in-flight sessions gets evicted from the session store and will fail with "Cosmos3 resume_session: unknown or expired session {session_id!r}." when it's next scheduled.
• Two concurrent requests can specify the same session_id, in which case both are writing to the same session (with potentially different data), and the last write wins.
• There's no TTL on the sessions, which is fine w.r.t. memory usage since states are small and it's a bounded queue, but does mean that a client that crashes mid-rollout leaves a resumable world state with no way to release it.

I loosely envision adding the following to /generate:

• start_session: bool
• session_timeout_s: int (with some default and maximum; we don't want sessions living on forever)
• session_id: str | None (if None and start_session is set, then the server creates a UUID for the session ID and returns it as an initial "special" result chunk.
• resume_session: bool (requires a session id)
• end_session: bool (requires a session id)

Then, the system can handle the resume session logic, persistent session state (e.g., having a sessions_states handled by the engine that mirrors the NodeSubmodule.request_states, giving Resources the option to retain state over a session instead of just at the request level), session timeout / cleanup, handling what happens if there are multiple in-flight requests with the same session (at least in the cosmos case, it seems like two requests with the same session is undefined behavior, but that might not be the case in the instance of bidirectional streaming).

@merceod I'm not sure if system-level abstractions for session state should be a blocker for this PR, but I can get started on it and put up a PR, since I have an idea of what it should look like.

[NSagan271]

Nit: it might be better organizationally for the cosmos3/tests files to all be in test/cosmos3 or test/modular/cosmos3 for consistency with other models

[merceod]

mstar/model/cosmos3/tests/ is where main already keeps the ten cosmos3 test files (test_action, test_droid, test_engine_cache, ...), so the new ones followed it. Happy to move all of them to test/cosmos3/ in a follow-up PR if you want one layout across models; doing it here would touch files this PR otherwise leaves alone.

[NSagan271]

Nit: If overlap_frames < tf/2, this could round down to zero, which I'd say is the expected behavior but maybe makes sense to emit a warning

[merceod]

Added a warning when a non-zero overlap_frames rounds to no latent frames (bc5a5d6), with a test.

[NSagan271]

This LRU ordering disagrees with the DiT (submodules.py:701), which promotes on read

[merceod]

Both nodes now use one SessionStore (a3c1d7c, ec768ca), and it promotes on read in both, so the two stores age their sessions the same way.

[NSagan271]

If there are 8 in-flight requests with live sessions and a 9th is added, then one of the in-flight sessions gets evicted from the session store and will fail with "Cosmos3 resume_session: unknown or expired session {session_id!r}." when it's next scheduled.

[merceod]

Fixed by pinning: a request pins its session at the prefill, pinned sessions are never evicted (only idle ones count against session_store_size), and the pin goes when the request stores its last window or is removed (cleanup_request). A second request on a live session is refused with "session 'x' is in use by request 'y'" instead of both writing it (a3c1d7c, tests in d4fa718).

[merceod]

Agreed on the split, and yes please, go ahead with the system-level session state; I'll move the cosmos3 store onto it when it lands. For now the model-side store handles the four points: one SessionStore for both nodes with promotion on read, sessions pinned while a request runs (never evicted, a second request on the same session_id is refused with a clear error, the pin goes with the request on completion or removal), a timeout on idle sessions (session_timeout_s, default 600 s, capped at 3600), and end_session to drop one early. Docs rows and CPU tests are in, plus a served stress for the same cases.

Also on this branch since your review: the KV retention now rides on the commit step's KVStep instead of set_retention (from the #293 thread), the overlap-rounding warning, and the fixes from #293-#296 carried by cherry-pick until those merge. The engine-cache harness passes 15/15 on the new retention path (windowed kv bit-identical to the oracle, pages released on the live request).

[NSagan271]

LGTM (pending the prerequisite branches being merged + a rebase0