mstar.model.cosmos3.submodules

NodeSubmodule wrappers for the Cosmos3 generator nodes.

Nodes:

Cosmos3DiTSubmodule – dual-pathway DiT (KV_CACHE). Dispatches by

graph_walk between prefill (the understanding tower runs once over the text prompt and writes its per-layer K/V) and image_gen (one denoising step of the generation tower per loop iteration, attending to the frozen understanding K/V plus the current generation tokens, then one scheduler step). Classifier-free guidance keeps the conditional and unconditional prompts in two cache labels and combines their velocities.

Cosmos3VAEEncoderSubmodule – Wan VAE conditioning encode (STATELESS): the

request’s conditioning image/video to clean anchor latents, in parallel with the DiT prefill.

Cosmos3VAEDecoderSubmodule – Wan VAE decode (STATELESS): final latents to

pixels.

Cosmos3VAEDecoderARSubmodule – streaming Wan VAE decode (STATELESS) for

windowed AR video: one committed window’s latents per stream chunk, decoded behind re-decoded context and emitted per window or assembled.

Cosmos3VisionEncoderSubmodule – Edge reasoner vision tower (STATELESS):

packed image/video patches to text-space tokens for the reasoner prefill.

Cosmos3ReasonerSubmodule – the understanding tower as a causal VLM

(shares the DiT’s transformer instance and kv/attn resources, plus its own sampler): reasoner_prefill / reasoner_prefill_vision write the prompt’s K/V and sample the first token, reasoner_decode one token per loop iteration.

Because the text tokens never receive a timestep embedding, the understanding K/V is denoise-step independent, so writing it once and re-reading it every step matches running the whole transformer each step.

Functions

native_flow_sigmas(num_inference_steps, ...)

The native flow-matching sigma grid: num_inference_steps values linearly spaced from 1 - 1/T toward 0, the endpoint dropped.

Classes

Cosmos3AudioDecoderSubmodule(...)	AVAE sound decode node (STATELESS): final denoised sound latents [1, C, T] -> stereo waveform [channels, samples] in [-1, 1], trimmed to the request's target sample count.
Cosmos3DiTSubmodule(transformer, config[, ...])	Dual-pathway DiT node (understanding tower + generation denoiser).
Cosmos3ReasonerSubmodule(transformer, config)	The understanding tower served as a causal VLM.
Cosmos3VAEDecoderARSubmodule(vae, config)	Streaming Wan VAE decode for windowed AR video (ported from #198).
Cosmos3VAEDecoderSubmodule(vae, config)	Wan VAE decode node: final denoised latents -> pixel frames.
Cosmos3VAEEncoderSubmodule(vae, config)	Wan VAE conditioning-encode node (STATELESS): the request's conditioning image or video -> clean anchor latents for the denoise loop.
Cosmos3VisionEncoderSubmodule(vision_model, ...)	The Edge reasoner's vision tower + projector (STATELESS): the request's packed image/video patches -> one text-space token per merged 2x2 block, in prompt order, for the reasoner prefill to scatter over its <|image_pad|> / <|video_pad|> tokens.
GenStepInfo(cfg, cfg_active, capture_key[, ...])	What a denoise step's declaration needs and cannot read off the batch.

class mstar.model.cosmos3.submodules.Cosmos3AudioDecoderSubmodule(sound_tokenizer, config)

Bases: NodeSubmodule

AVAE sound decode node (STATELESS): final denoised sound latents [1, C, T] -> stereo waveform [channels, samples] in [-1, 1], trimmed to the request’s target sample count.

The target is re-derived from the request metadata (duration x sample rate) rather than read from the DiT node’s per-request state, so this node stays stateless and placeable on any rank.

disable_torch_compile: bool = True

forward(graph_walk, engine_inputs, sound_latents, target_samples, **kwargs)

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:

engine_inputs (ModelInputsFromEngine)

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)

Return type:

NodeInputs

class mstar.model.cosmos3.submodules.Cosmos3DiTSubmodule(transformer, config, scheduler=None)

Bases: ARNodeSubmodule

Dual-pathway DiT node (understanding tower + generation denoiser).

batched_cfg = True

can_batch(batch, model_inputs)

Return type:

bool

cg_key_info(graph_walk, per_request_info, per_request_input_metadata=None, **kwargs)

Which of this walk’s capture buckets the batch belongs to, or None for “run eager”.

The engine leases the replay slot before the step is declared, so this answers from per-request state rather than from prepared inputs. It is the sole gate on capture — the v1 engine has no can_use_cuda_graphs — so every condition the old gate checked lives here:

• only the two-branch guidance regime is captured (both the prefill’s combined cond+uncond pack and the denoise step’s batched CFG);

• of the denoise walks only image_gen, and only at a resolution a graph was captured for (_capture_layout);

• a batch shares one captured (batch size, token count) bucket, so a mixed-resolution batch falls back to the eager cross-request denoise.

Every fact this reads is fixed for a request’s whole lifetime (its guidance regime and its latent shape, both settled at prefill). That is a requirement, not a coincidence: the speculative pre-plan path leases a slot for step N+1 while step N is still in flight, before N+1 has been through prepare_inputs, so a key derived from anything that advances per step (the denoise step index, say) would be read one iteration stale and could lease a bucket the declaration then contradicts. Engine.exec checks the two against each other.

The guidance interval is the fact this deliberately does not consult: a captured step always runs both branches, so an out-of-interval step gets the guidance combine where the eager path would run the conditional branch alone (postprocess combines unconditionally). That predates the v1 migration and matches the fused reference pipeline, which applies guidance on every step; honoring the interval here instead would make the key step-dependent, which is exactly what the pre-plan path cannot support.

Parameters:

• graph_walk (str)

• per_request_info (dict)

Return type:

object | None

check_stop(request_id, request_info, outputs)

Stop this request’s denoise loop once it has run its own step count.

The loop is built with a fixed upper-bound iteration count (config.max_inference_steps); each request runs only as many steps as its scheduler holds (e.g. image 50, video 35, action 30, distilled policy ~4), which can differ between concurrent requests. Runs on the worker’s slow-postprocess path, so reading the per-request step count is fine. The one extra step the loop dispatches before this stop takes effect is vetoed by prepare_inputs (the forwards also guard it for direct callers, e.g. tests).

Return type:

set[str]

cleanup_request(request_id)

Remove per-request state when a request completes. The engines call this on request removal; overrides with extra internal state should call super().

Parameters:

request_id (str)

declare_step(graph_walk, request_ids, inputs, slot_lease=None, piecewise_leases=None, **kwargs)

This batch’s step: which cache streams it touches, by how much, and which attention backend runs over them.

Two guidance branches are two labels on the one cache. Whether a request has both, and whether this step uses both, is a per-request fact that rides in on NodeInputs.resource_step_info from prepare_inputs (see GenStepInfo) rather than being read off the batch here.

The attention key is the other half of the declaration. attn is paged, attn_gen dense; a step that may replay from a captured graph must name the paged one, because the dense backend’s prefix gather is shaped by the step and cannot be captured.

Parameters:

• graph_walk (str)

• request_ids (list[str])

• inputs (list[ARNodeInputs])

• slot_lease (SlotLease | None)

• piecewise_leases (Mapping[str, SlotLease] | None)

Return type:

SubmoduleStep

disable_torch_compile: bool = True

forward(graph_walk, engine_inputs, **kwargs)

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:

engine_inputs (ModelInputsFromEngine)

forward_batched(graph_walk, engine_inputs, latents=None, time_index=None, action_latents=None, sound_latents=None, input_ids=None, text_mrope_ids=None, **kwargs)

Batched form of forward: a multi-request batch’s outputs, per request or as a BatchedModelOutput. Override when can_batch returns True.

Parameters:

engine_inputs (ModelInputsFromEngine)

forward_captured(graph_walk, engine_inputs, latents, vision_timesteps, position_ids_cond, position_ids_uncond, noisy_token_mask, **kwargs)

Velocity-only denoise forward captured into a CUDA graph: both guidance branches in one pass (the combined plan), no scheduler step. The token layout is baked per resolution; the latents, timestep and rotary positions are static-buffer inputs stacked on a leading batch dim. A single request keeps the two-branch path; a concurrent batch runs the per-request denoise (the same compute as the eager cross-request forward), one transformer pass over the whole batch.

Parameters:

engine_inputs (ModelInputsFromEngine)

Return type:

dict

gen_capture_batch_sizes: tuple[int, ...] = (1,)

gen_capture_resolutions: tuple[tuple[int, int], ...] = ((192, 320), (480, 832), (720, 1280))

get_cuda_graph_configs(device, tp_world_size=1)

Declare one fixed-shape capture of the image denoise step per resolution. Requests at other resolutions, or without guidance, fall back to the eager path. The per-resolution token layout is prompt-independent, so bake it once here and key it by latent shape; the per-prompt rotary positions, the latents and the timestep flow in as static-buffer inputs.

Set COSMOS3_DISABLE_CUDA_GRAPH=1 to skip capture and run the denoise loop eagerly (escape hatch for a misbehaving driver, and an A/B switch). Set COSMOS3_GEN_CAPTURE_RES (e.g. "192x320,480x832", height x width) to override which resolutions are captured, and COSMOS3_GEN_CAPTURE_BS (e.g. "1,4,8") to also capture batched denoise steps so concurrent requests replay a padded graph instead of falling back to the eager path.

Parameters:

tp_world_size (int)

max_batch_size(graph_walk)

Parameters:

graph_walk (str)

max_gen_batch_size = 8

postprocess(request_id, request_info, outputs, inputs=None, **kwargs)

Captured-path tail: the classifier-free-guidance combine and the (Python, multistep) scheduler step the graph can’t hold, finished from the step’s own inputs. Mirrors the tail of _forward_image_gen. Eager forwards already emit finished latents/time_index and pass through untouched (no cond_v in their outputs).

prefill_capture_batch_sizes: tuple[int, ...] = (1,)

prefill_capture_token_buckets: tuple[int, ...] = (16, 32, 64, 128, 256)

prepare_inputs(graph_walk, fwd_info, inputs, seen_token_mask=None, pos_info={}, **kwargs)

Return type:

ARNodeInputs

preprocess(graph_walk, engine_inputs, inputs)

Parameters:

• engine_inputs (ModelInputsFromEngine)

• inputs (list[ARNodeInputs])

Return type:

dict

to(*args, **kwargs)

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

• device (torch.device) – the desired device of the parameters and buffers in this module

• dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module

• tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

• memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

class mstar.model.cosmos3.submodules.Cosmos3ReasonerSubmodule(transformer, config)

Bases: ARNodeSubmodule

The understanding tower served as a causal VLM.

Shares the DiT node’s Cosmos3OmniTransformer instance (one copy of the text weights) and its kv / attn resources; declares its own sampler. The prefill walks embed the rendered prompt, scatter the vision encoder’s tokens over the media placeholders, run the text tower (writing the raw K/V under one label), and sample the first token from the last position; the decode loop feeds each sampled token back as the next step’s single-token input. Positions are the prompt’s 3D mRoPE ids (computed in process_prompt) and, past the prompt, the scalar cursor max(position) + 1 on all three axes, carried in per-request state.

can_batch(batch, model_inputs)

Return type:

bool

check_stop(request_id, request_info, outputs)

Return the set of dynamic-loop names that should stop after this step.

Runs on the worker’s slow-postprocess path after execute_batch returns — never inside execute_batch. Allowed to read tensor values (.item() / .cpu()) because by this point the GPU thread is no longer blocked by it.

Stops returned here are deferred by one step: they apply to the worker’s next iter’s fast postprocess. The current in-flight step (already submitted under the assumption that the rid continues) will run for that rid and its output discarded — the standard 1-wasted-step cost for any stop signal.

Default: no stops.

Return type:

set[str]

declare_step(graph_walk, request_ids, inputs, slot_lease=None, piecewise_leases=None, **kwargs)

One causal span per request under the reasoner’s label, committed (the text is context for every later token); the sampler tracks the prompt tokens at prefill for the repetition penalty.

Parameters:

• graph_walk (str)

• request_ids (list[str])

• inputs (list[ARNodeInputs])

• slot_lease (SlotLease | None)

• piecewise_leases (Mapping[str, SlotLease] | None)

Return type:

SubmoduleStep

decode_capture_batch_sizes: tuple[int, ...] = (1, 2, 4, 8, 16, 32)

disable_torch_compile: bool = True

forward(graph_walk, engine_inputs, input_ids, position_ids, seq_lens, vision_embeds=None, **kwargs)

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:

engine_inputs (ModelInputsFromEngine)

forward_batched(graph_walk, engine_inputs, input_ids, position_ids, seq_lens, vision_embeds=None, **kwargs)

Batched form of forward: a multi-request batch’s outputs, per request or as a BatchedModelOutput. Override when can_batch returns True.

Parameters:

engine_inputs (ModelInputsFromEngine)

get_cuda_graph_configs(device, tp_world_size=1)

Capture the decode step (one token per request) per batch-size bucket; prefills run eager (they are one-shot and shape-varied).

Parameters:

tp_world_size (int)

postprocess(request_id, request_info, outputs, inputs=None, **kwargs)

Per-request postprocessing on the submodule outputs.

Runs on the GPU thread inside execute_batch, after the forward (eager, batched, or CUDA-graph replay). inputs is the request’s prepare_inputs result for this step, so a submodule can finish a step the captured graph could not hold (e.g. combine guidance branches and run a Python multistep scheduler against the step’s input latents).

Keep it metadata-only where possible. Avoid reading tensor values — .item() / .cpu() sync here block the GPU thread and forfeit the worker’s async-scheduling overlap; stop-condition decisions that need token values (e.g. EOS) belong in check_stop. A captured-path tail that must read scheduler state is the sanctioned exception — it costs the same sync wherever it runs.

Typical uses:

• rebind output names for graph routing (outputs["text_inputs"] = outputs["new_token"]);

• drop keys on a per-request basis for static-capture submodules (e.g. Qwen3-Omni Thinker dropping thinker_states for requests that don’t need audio);

• finish a captured step from inputs (Cosmos3 denoise tail).

Modifies outputs in-place; returns nothing.

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)

Return type:

ARNodeInputs

preprocess(graph_walk, engine_inputs, inputs)

Parameters:

• engine_inputs (ModelInputsFromEngine)

• inputs (list[ARNodeInputs])

Return type:

dict

class mstar.model.cosmos3.submodules.Cosmos3VAEDecoderARSubmodule(vae, config)

Bases: Cosmos3VAEDecoderSubmodule

Streaming Wan VAE decode for windowed AR video (ported from #198).

Consumes one committed window’s latents per stream chunk. Each window is decoded behind a re-decoded left context (the last few latents of the stream so far) so the causal conv stack is warm at the kept frames; the context- and overlap-derived pixels are trimmed, and the video is either emitted per window (stream_video) or assembled and emitted once the last window lands. The chunk count is the completion signal — it is known per request up front — so the stream’s terminal flush (an empty pass, and the only pass a non-windowed request’s idle stream ever delivers) runs as a no-op forward: the pass must complete normally for the partition to report done, so it is not vetoed.

forward(graph_walk, engine_inputs, latents, **kwargs)

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:

engine_inputs (ModelInputsFromEngine)

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)

Return type:

NodeInputs

class mstar.model.cosmos3.submodules.Cosmos3VAEDecoderSubmodule(vae, config)

Bases: NodeSubmodule

Wan VAE decode node: final denoised latents -> pixel frames.

Applies the pipeline-side latent normalization (the VAE itself returns raw latents) before decoding, matching the fused t2i pipeline’s decode.

disable_torch_compile: bool = True

forward(graph_walk, engine_inputs, latents, **kwargs)

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:

engine_inputs (ModelInputsFromEngine)

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)

Return type:

NodeInputs

class mstar.model.cosmos3.submodules.Cosmos3VAEEncoderSubmodule(vae, config)

Bases: NodeSubmodule

Wan VAE conditioning-encode node (STATELESS): the request’s conditioning image or video -> clean anchor latents for the denoise loop.

Runs in parallel with the DiT’s understanding-tower prefill (the conditioned prefill walks put the two nodes in a Parallel section) and emits cond_latents as a persist signal the conductor threads into the generation walk’s first iteration. Everything is re-derived from the request metadata so the node stays stateless and placeable on any rank. Per mode:

• image-to-video: the single conditioning frame, encoded standalone (the Wan VAE is temporally causal, so frame 0 encodes bit-identically alone);

• action policy / forward-dynamics: the frame repeated across the clip, full-clip encoded (mirrors the reference pipelines);

• action inverse-dynamics: the whole observed clip;

• video-to-video: the condition_video_keep prefix (short clips padded by repeating the last frame), encoded and scattered into the pinned frames of a full-latent-shape tensor (zeros elsewhere).

disable_torch_compile: bool = True

forward(graph_walk, engine_inputs, vision, condition_indexes=None, latent_shape=None, **kwargs)

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:

engine_inputs (ModelInputsFromEngine)

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)

Shape the conditioning pixels for the encode: mode-dependent frame selection/padding plus the [-1, 1] resize-normalize, all derived from the request metadata (no per-request state).

Return type:

NodeInputs

class mstar.model.cosmos3.submodules.Cosmos3VisionEncoderSubmodule(vision_model, config)

Bases: NodeSubmodule

The Edge reasoner’s vision tower + projector (STATELESS): the request’s packed image/video patches -> one text-space token per merged 2x2 block, in prompt order, for the reasoner prefill to scatter over its <|image_pad|> / <|video_pad|> tokens.

The pixel patches and their grids are computed CPU-side in Cosmos3Model.process_prompt (the token count must be known when the prompt is rendered), so this node only runs the encoder.

disable_torch_compile: bool = True

forward(graph_walk, engine_inputs, pixel_values, vision_grid_thw, **kwargs)

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:

engine_inputs (ModelInputsFromEngine)

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)

Return type:

NodeInputs

class mstar.model.cosmos3.submodules.GenStepInfo(cfg, cfg_active, capture_key, commit=False, retention=None)

Bases: object

What a denoise step’s declaration needs and cannot read off the batch.

All three are per-request facts resolved in prepare_inputs (which has the request’s state and its step index) and carried to declare_step on NodeInputs.resource_step_info, since the declaration runs with only the request ids and their prepared inputs.

capture_key is this request’s half of cg_key_info — the capture bucket its shape belongs to, or None for “no captured graph for this”. Agreeing with cg_key_info matters: the engine leases the slot from that and then checks the declaration against the lease. What is left over is the batch’s size and token count, which no row knows, so a step every row of which is capturable may still find no bucket and run eager. Declaring the paged backend for such a step only costs the dense path’s speedup; declaring the dense one for a step that does get a slot would try to capture an eager-only backend, so the safe direction over-declares.

cfg_active narrows the eager single-request declaration only. It is deliberately not part of capture_key — see cg_key_info.

Parameters:

• cfg (bool)

• cfg_active (bool)

• capture_key (object | None)

• commit (bool)

• retention (Mapping[str, RetentionPolicy] | None)

capture_key: object | None

cfg: bool

cfg_active: bool

commit: bool = False

retention: Mapping[str, RetentionPolicy] | None = None

mstar.model.cosmos3.submodules.native_flow_sigmas(num_inference_steps, num_train_timesteps)

The native flow-matching sigma grid: num_inference_steps values linearly spaced from 1 - 1/T toward 0, the endpoint dropped. A numpy array: UniPC’s explicit-sigma path applies the flow shift arithmetically to it (a Python list raises inside set_timesteps).

Parameters:

• num_inference_steps (int)

• num_train_timesteps (int)