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.

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.

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).
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.
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)

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]

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: maps a multi-request batch to per-request outputs. 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, **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.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.GenStepInfo(cfg, cfg_active, capture_key)

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)

capture_key: object | None

cfg: bool

cfg_active: bool