mstar.model.submodule_base

Classes

ARNodeInputs([tensor_inputs, kwargs, ...])	Unlike in regular ModelInputs, for LLMInputs we expect either input_ids or input_embeds to be set (but typically not both), and we require input_seq_len to be set (for cache planning).
ARNodeSubmodule()
BatchedModelOutput([per_rid_outputs, ...])	A step's outputs, split by how they are addressed.
HostRows(request_ids, buffers)	A step's row-addressed outputs (check_stop_buffers) on the host.
InputMetadata([stream_chunks, ...])	What the worker knows about one request's inputs to this step.
LazyRequestStates(submodule, rids)	The batch's ``PerRequestState``s, resolved on first read.
ModelInputsFromEngine(request_ids, ...)
NodeInputs(tensor_inputs, kwargs, ...)
NodeSubmodule()	Base class for a model's compute units: defines the prepare_inputs → preprocess → forward(_batched) contract the engines drive.
PerRequestState([tensors, kwargs, ...])	Engine-owned per-request state a submodule persists across forwards.
StackingMethod(*values)

class mstar.model.submodule_base.ARNodeInputs(tensor_inputs=<factory>, kwargs=<factory>, resource_step_info=None, input_seq_len=0, input_ids=None, input_embeds=None, custom_pos_ids=None)

Bases: NodeInputs

Unlike in regular ModelInputs, for LLMInputs we expect either input_ids or input_embeds to be set (but typically not both), and we require input_seq_len to be set (for cache planning).

The tensor_inputs and kwargs dicts are still available for additional inputs as needed; but the main LLM inputs should be provided in the given dedicated fields.

Parameters:

• tensor_inputs (dict[str, Tensor])

• kwargs (dict)

• resource_step_info (Any | None)

• input_seq_len (int)

• input_ids (Tensor | None)

• input_embeds (Tensor | None)

• custom_pos_ids (Tensor | dict[str, Tensor] | None)

classmethod collate(inputs_list, stacking_method=StackingMethod.NONE)

Parameters:

inputs_list (list[ARNodeInputs])

custom_pos_ids: Tensor | dict[str, Tensor] | None = None

input_embeds: Tensor | None = None

input_ids: Tensor | None = None

class mstar.model.submodule_base.ARNodeSubmodule

Bases: NodeSubmodule

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

Parameters:

• graph_walk (str)

• fwd_info (CurrentForwardPassInfo)

• inputs (dict[str, list[Tensor]])

Return type:

ARNodeInputs

abstractmethod preprocess(graph_walk, engine_inputs, inputs)

Parameters:

• graph_walk (str)

• engine_inputs (ModelInputsFromEngine)

• inputs (list[ARNodeInputs])

Return type:

dict[str, Tensor | Any]

split_inputs(graph_walk, fwd_info, inputs, start, end)

This walk’s inputs restricted to tokens [start, end).

Only the sequence-shaped fields are cut. tensor_inputs, kwargs and resource_step_info are opaque here, so a submodule carrying any of them says for itself how they are cut, or refuses to be cut at all.

Parameters:

• graph_walk (str)

• fwd_info (CurrentForwardPassInfo)

• inputs (ARNodeInputs)

• start (int)

• end (int)

Return type:

ARNodeInputs

class mstar.model.submodule_base.BatchedModelOutput(per_rid_outputs=<factory>, packed_outputs=<factory>, check_stop_buffers=None, row_outputs=None, row_request_ids=None)

Bases: object

A step’s outputs, split by how they are addressed.

per_rid_outputs is keyed by request id. packed_outputs holds whole batch tensors a captured graph emitted under a __name__ key, which a submodule’s unpack_packed_outputs cuts per request.

check_stop_buffers is what the stop check reads (default: the per-rid outputs, one device-to-host copy per tensor per request); a batch tensor with row i for request i makes it one copy.

row_outputs does the same for the outputs: the engine clones each batch tensor once, not once per request, and hands out row views. Only Engine._collect_outputs reads it.

The two dicts are never None, so readers on the critical path need no guard.

Parameters:

• per_rid_outputs (dict[str, dict[str, list[Tensor]]])

• packed_outputs (dict[str, Tensor])

• check_stop_buffers (dict[str, Tensor | dict[str, list[Tensor]]] | None)

• row_outputs (dict[str, Tensor] | None)

• row_request_ids (tuple[str, ...] | None)

check_stop_buffers: dict[str, Tensor | dict[str, list[Tensor]]] | None = None

clone_check_stop_buffers()

A detached copy: the next replay overwrites a captured graph’s buffers before the stop check reads them.

clone_row_outputs(num_rows)

One clone per row-addressed batch tensor, narrowed to the real rows so padding rows are never copied (see clone_check_stop_buffers).

Parameters:

num_rows (int)

Return type:

dict[str, Tensor]

classmethod coerce(output)

Parameters:

output (BatchedModelOutput | dict[str, Any])

Return type:

BatchedModelOutput

get(key, default=None)

Parameters:

key (str)

get_check_stop_input()

packed_outputs: dict[str, Tensor]

per_rid_outputs: dict[str, dict[str, list[Tensor]]]

pop(request_id, default=None)

Drop one request’s per-rid outputs; row-addressed tensors are left alone, and the caller takes the request out of the batch instead.

Parameters:

request_id (str)

row_outputs: dict[str, Tensor] | None = None

row_request_ids: tuple[str, ...] | None = None

update(other)

Parameters:

other (BatchedModelOutput | dict[str, Any])

class mstar.model.submodule_base.HostRows(request_ids, buffers)

Bases: NamedTuple

A step’s row-addressed outputs (check_stop_buffers) on the host.

Row i of every buffer belongs to request_ids[i], the order the forward ran the requests in; a padded replay leaves extra rows past the real ones.

Parameters:

• request_ids (tuple)

• buffers (dict[str, Tensor])

buffers: dict[str, Tensor]

Alias for field number 1

request_ids: tuple

Alias for field number 0

class mstar.model.submodule_base.InputMetadata(stream_chunks=mappingproxy({}), dynamic_loop_iter_counts=mappingproxy({}))

Bases: NamedTuple

What the worker knows about one request’s inputs to this step.

Per step, unlike the per-request CurrentForwardPassInfo. Build it with keywords: fields get added here rather than as new hook arguments.

Parameters:

• stream_chunks (Mapping[str, StreamChunkInfo])

• dynamic_loop_iter_counts (Mapping[str, int])

dynamic_loop_iter_counts: Mapping[str, int]

Alias for field number 1

stream_chunks: Mapping[str, StreamChunkInfo]

Alias for field number 0

class mstar.model.submodule_base.LazyRequestStates(submodule, rids)

Bases: Mapping

The batch’s ``PerRequestState``s, resolved on first read.

Only a couple of submodules read these, but the engine builds the view for every step of every node; materialising the dict there is bs dict inserts (and, for a padded step, bs PerRequestState allocations) per step that nothing usually looks at.

Parameters:

• submodule (NodeSubmodule)

• rids (Sequence[str])

class mstar.model.submodule_base.ModelInputsFromEngine(request_ids: 'list[str]', per_request_info: 'dict[str, CurrentForwardPassInfo]', resources: 'dict[str, Resource]'=<factory>, piecewise_runners: "dict[str, 'PiecewiseCudaGraphRunner']"=<factory>, per_request_states: "'Mapping[str, PerRequestState] | None'"=None, step: 'SubmoduleStep | None' = None, captured: 'bool' = False, per_request_input_metadata: 'Mapping[str, InputMetadata]'=<factory>)

Bases: object

Parameters:

• request_ids (list[str])

• per_request_info (dict[str, CurrentForwardPassInfo])

• resources (dict[str, Resource])

• piecewise_runners (dict[str, 'PiecewiseCudaGraphRunner'])

• per_request_states (Mapping[str, PerRequestState] | None)

• step (SubmoduleStep | None)

• captured (bool)

• per_request_input_metadata (Mapping[str, InputMetadata])

captured: bool = False

property first_request_info

unlike single_request_info, does not assert that there is only one request

per_request_info: dict[str, CurrentForwardPassInfo]

per_request_input_metadata: Mapping[str, InputMetadata]

per_request_states: Mapping[str, PerRequestState] | None = None

piecewise_runners: dict[str, 'PiecewiseCudaGraphRunner']

request_ids: list[str]

resources: dict[str, Resource]

property single_request_info

asserts that there is only one request

Type:

IMPORTANT

step: SubmoduleStep | None = None

class mstar.model.submodule_base.NodeInputs(tensor_inputs: 'dict[str, torch.Tensor]'=<factory>, kwargs: 'dict' = <factory>, resource_step_info: 'Any | None' = None, input_seq_len: 'int' = 0)

Bases: object

Parameters:

• tensor_inputs (dict[str, Tensor])

• kwargs (dict)

• resource_step_info (Any | None)

• input_seq_len (int)

clone()

Copy with tensors cloned, so a capture template can be reused.

Goes through the fields rather than naming them, so a subclass gets its own type back without restating this.

input_seq_len: int = 0

kwargs: dict

resource_step_info: Any | None = None

tensor_inputs: dict[str, Tensor]

class mstar.model.submodule_base.NodeSubmodule

Bases: Module, ABC

Base class for a model’s compute units: defines the prepare_inputs → preprocess → forward(_batched) contract the engines drive.

bind_node_resources(resources)

Receive the engine-built resources for this submodule’s node, and pass them down to every layer that calls one.

A layer body (attention, cross-attention) calls the resources directly — attn.run, kv.write_kv, pos.apply_qk — so it needs its own references. It resolves them here, once at load, by the labels the model declared in get_node_resources; a layer that names a label this node doesn’t have fails at bind rather than in the middle of a forward.

Parameters:

resources (dict[str, Any])

Return type:

None

can_batch(batch, model_inputs)

Parameters:

• batch (ExecutingBatch)

• model_inputs (list[NodeInputs])

can_use_cuda_graphs(batch, model_inputs)

Return True if this submodule supports CUDA graphs for batch.

Default: derives from get_cuda_graph_configs — if any declared config can replay for this batch’s graph_walk, CUDA graphs are supported. We check cfg.replay_graph_walks (not just cfg.capture_graph_walk) so aliased walks — e.g. Qwen3-Omni’s prefill_audio reusing the prefill_text capture, or prefill_vision reusing its own — are correctly admitted at the eligibility gate. The runner’s _config_for already looks up by replay_graph_walks; this keeps the gate consistent so aliased walks don’t silently fall through to the eager path.

replay_graph_walks is always a superset of {capture_graph_walk} (see CudaGraphConfig.__init__), so this never narrows what the previous code accepted — only widens it for configs that explicitly declared aliases.

Subclasses can override to reject on batch shape / metadata (e.g. codec submodules that need homogeneous frame counts).

Parameters:

• batch (ExecutingBatch)

• model_inputs (list[NodeInputs])

Return type:

bool

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

Which of this walk’s capture buckets a batch belongs to.

A walk can be captured more than once when the batch’s shape is not the whole story — bagel captures decode twice, guidance on and off, because the two declare different segments over the same token count. The engine leases a slot before the step is declared, so it cannot read the answer off the step; it asks here instead.

Must equal the additional_key_info on the config that captured the bucket, and the cg_key_info this batch’s declare_step puts on its step. Disagreeing is not an error anywhere — it just misses the capture and runs eager — so derive both from one place.

per_request_input_metadata is the batch’s per-step input facts (as on ModelInputsFromEngine), e.g. for a bucket that follows a stream chunk.

None (the default) means the walk has a single capture.

Parameters:

• graph_walk (str)

• per_request_info (dict[str, CurrentForwardPassInfo])

• per_request_input_metadata (Mapping[str, InputMetadata] | None)

• kwargs (Any)

Return type:

Any

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.

Parameters:

• request_id (str)

• request_info (CurrentForwardPassInfo)

• outputs (dict[str, list[Tensor]])

Return type:

set[str]

check_stop_batched(request_ids, request_infos, host_rows)

check_stop for the whole batch at once, off the host rows.

Optional. Returns rid -> loops to stop (rids with none left out), and must agree with check_stop for every rid. None, or a raise, falls back to the per-request path.

Parameters:

• request_ids (list[str])

• request_infos (dict[str, CurrentForwardPassInfo])

• host_rows (HostRows)

Return type:

dict[str, set[str]] | None

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)

Declare this batch’s step for the runner to drive: which cache streams it touches, what spans they grow by, which plans back it, which streams fork, and what commits when it lands. The runner drives the declaration before preprocess and commits it after the forward, so a declaring submodule keeps no plan or advance calls of its own. None means the submodule still plans and advances through the facade itself.

request_ids pairs positionally with inputs. Under a captured graph the batch is padded to the bucket’s shape, so it carries the padding rows’ ids too — declare their segments like any other row.

slot_lease is the slot this step will replay on, or None for an eager step. A submodule whose declaration differs between the two (cosmos3 packs both guidance branches into one plan for the captured shape) must key off this, not off its own capture key: the key says the batch could be captured, the lease says it was.

piecewise_leases names the regions of this node that hold a slot for this step. Such a region declares, plans and commits its own work, so a resource it owns must be left out of this declaration.

Parameters:

• graph_walk (str)

• request_ids (list[str])

• inputs (list[NodeInputs])

• slot_lease (SlotLease | None)

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

Return type:

SubmoduleStep | None

disable_autocast: bool = False

disable_torch_compile: bool = False

filter_batched_output(request_info, outputs)

Drop keys a real request shouldn’t receive. A captured forward emits a fixed key set for graph compat, so the filtering happens here.

Parameters:

• request_info (CurrentForwardPassInfo)

• outputs (dict[str, list[Tensor]])

Return type:

dict[str, list[Tensor]]

abstractmethod forward(graph_walk, engine_inputs, **kwargs)

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

Parameters:

• graph_walk (str)

• engine_inputs (ModelInputsFromEngine)

Return type:

dict[str, list[Tensor]]

forward_batched(graph_walk, engine_inputs, **kwargs)

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

Parameters:

• graph_walk (str)

• engine_inputs (ModelInputsFromEngine)

Return type:

dict[str, dict[str, list[Tensor]]] | BatchedModelOutput

get_autocast_dtype()

Per-submodule autocast dtype override for the engine’s forward wrap. The engine consults this on each execute_batch and uses the returned dtype instead of its own when non-None.

Default: None (inherit the engine’s autocast dtype). To turn autocast off for one specific submodule whose engine otherwise has it enabled, wrap the submodule’s forward with torch.amp.autocast(enabled=False) — that path is engine-agnostic and doesn’t need this surface.

Return type:

dtype | None

get_cuda_graph_configs(device, tp_world_size=1)

Parameters:

• device (torch.device)

• tp_world_size (int)

Return type:

list[CudaGraphConfig]

get_device()

get_piecewise_cuda_graph_configs(device, autocast_dtype, tp_world_size=1)

Return the piecewise CUDA graph configs this submodule opts into.

autocast_dtype is the engine’s autocast dtype — passed so a config’s make_static_inputs can allocate the hidden-state buffer in the dtype the captured region runs under (avoids a copy-time upcast at replay).

A piecewise CUDA graph captures ONE inner callable of this submodule’s forward (e.g. a transformer block loop) as a CUDA graph while the surrounding compute stays eager. The engine builds one PiecewiseCudaGraphRunner per returned label and threads the runners into ModelInputsFromEngine.piecewise_runners so the submodule’s forward can look them up by label:

runner = engine_inputs.piecewise_runners.get(“block_loop”) if runner is not None and runner.can_run(bs):

out = runner.run(static_inputs={…}, request_ids=…, seq_lens=…)

A config’s capture_fn takes one PiecewiseCallInputs, whose engine_inputs.resources carries the node’s resources — so a resource call (sampling included) can live INSIDE the capture, reading params straight from buffers whose addresses are stable across replays instead of being hoisted out as static inputs. A region that touches a resource declares its work in the config’s own declare_step; the runner admits, plans and commits it per replay.

Default: no piecewise graphs. Override to return {label: PiecewiseCudaGraphConfig}; multiple labels capture multiple independent graphs (i.e., one per outer function to be graphed).

Parameters:

• device (torch.device)

• autocast_dtype (torch.dtype)

• tp_world_size (int)

Return type:

dict[str, PiecewiseCudaGraphConfig]

max_batch_size(graph_walk)

Parameters:

graph_walk (str)

node_resources: dict[str, Any]

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.

Parameters:

• request_id (str)

• request_info (CurrentForwardPassInfo)

• outputs (dict[str, list[Tensor]])

• inputs (NodeInputs | None)

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

Parameters:

• graph_walk (str)

• fwd_info (CurrentForwardPassInfo)

• inputs (dict[str, list[Tensor]])

Return type:

NodeInputs

preprocess(graph_walk, engine_inputs, inputs)

Parameters:

• graph_walk (str)

• engine_inputs (ModelInputsFromEngine)

• inputs (list[NodeInputs])

Return type:

dict[str, Tensor | Any]

request_state(request_id)

The request’s state, created on first access.

Parameters:

request_id (str)

Return type:

PerRequestState

request_states: dict[str, PerRequestState]

split_batches_by_capture_key(graph_walk)

Whether the scheduler batches each cg_key_info group apart.

A batch spanning two groups matches no capture. Opt in when that batch runs worse than two captured ones, e.g. per-request eager; leave it off when a batched eager forward handles the mix well.

Parameters:

graph_walk (str)

Return type:

bool

split_inputs(graph_walk, fwd_info, inputs, start, end)

This walk’s inputs restricted to tokens [start, end).

inputs is the whole prompt, prepared once and sliced rather than re-derived, so anything read out of a resource while preparing is read once, before any part of it runs.

Parameters:

• graph_walk (str)

• fwd_info (CurrentForwardPassInfo)

• inputs (NodeInputs)

• start (int)

• end (int)

Return type:

NodeInputs

unpack_packed_outputs(static_output, request_ids, real_seq_lens, inputs, per_request_info)

Per-rid slicing for packed sentinels emitted by the captured graph.

Decode-style submodules emit per-rid entries inside the captured forward (one slice per request, fixed shape), so they don’t need this. Prefill-style submodules pack a (total_tokens, …) tensor whose per-request slice ends depend on real seq_lens — slicing has to happen post-replay, outside the captured region. Default no-ops; override and key off static_output sentinel names.

Parameters:

• static_output (dict)

• request_ids (list[str])

• real_seq_lens (list[int])

• inputs (list[NodeInputs])

• per_request_info (dict[str, CurrentForwardPassInfo])

Return type:

dict[str, dict[str, list[Tensor]]]

class mstar.model.submodule_base.PerRequestState(tensors=<factory>, kwargs=<factory>, disag_shared_keys=<factory>)

Bases: object

Engine-owned per-request state a submodule persists across forwards.

Submodules stash whatever a request’s later steps need (schedulers, conditioning latents, packing metadata) instead of keeping private dict[request_id, ...] attributes. The engine owns the lifecycle: it injects the batch’s states via ModelInputsFromEngine.per_request_states and drops a request’s state when the request is removed — no submodule cleanup code required.

tensors vs kwargs split by value kind: device tensors go in tensors, everything else (numbers, dicts, scheduler objects) in kwargs. disag_shared_keys is reserved for PD disaggregation — marked keys would travel with the request (tensors via the tensor manager, kwargs with the forward-pass info); no engine implements the transfer yet.

Parameters:

• tensors (dict[str, Tensor])

• kwargs (dict[str, Any])

• disag_shared_keys (set[str])

add(key, value)

Parameters:

key (str)

Return type:

None

add_all(**kwargs)

Return type:

None

disag_shared_keys: set[str]

get(key, default=None)

Parameters:

key (str)

kwargs: dict[str, Any]

remove(keys)

Return type:

None

tensors: dict[str, Tensor]

class mstar.model.submodule_base.StackingMethod(*values)

Bases: Enum

CAT = 'cat'

NONE = 'none'

STACK = 'stack'