mstar.model.submodule_base

Contents

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

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)[source]#

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:
classmethod collate(inputs_list, stacking_method=StackingMethod.NONE)[source]#
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[source]#

Bases: NodeSubmodule

abstractmethod prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#
Parameters:
Return type:

ARNodeInputs

abstractmethod preprocess(graph_walk, engine_inputs, inputs)[source]#
Parameters:
Return type:

dict[str, Tensor | Any]

class mstar.model.submodule_base.LazyRequestStates(submodule, rids)[source]#

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:
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)[source]#

Bases: object

Parameters:
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_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)[source]#

Bases: object

Parameters:
clone()[source]#

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[source]#

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)[source]#

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)[source]#
Parameters:
can_use_cuda_graphs(batch, model_inputs)[source]#

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:
Return type:

bool

cg_key_info(graph_walk, per_request_info)[source]#

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.

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

Parameters:
Return type:

Any

check_stop(request_id, request_info, outputs)[source]#

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:
Return type:

set[str]

cleanup_request(request_id)[source]#

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)[source]#

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:
Return type:

SubmoduleStep | None

disable_autocast: bool = False#
disable_torch_compile: bool = False#
filter_batched_output(request_info, outputs)[source]#

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:
Return type:

dict[str, list[Tensor]]

abstractmethod forward(graph_walk, engine_inputs, **kwargs)[source]#

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

Parameters:
Return type:

dict[str, list[Tensor]]

forward_batched(graph_walk, engine_inputs, **kwargs)[source]#

Batched form of forward: maps a multi-request batch to per-request outputs. Override when can_batch returns True.

Parameters:
Return type:

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

get_autocast_dtype()[source]#

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)[source]#
Parameters:
Return type:

list[CudaGraphConfig]

get_device()[source]#
get_piecewise_cuda_graph_configs(device, autocast_dtype, tp_world_size=1)[source]#

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:
Return type:

dict[str, PiecewiseCudaGraphConfig]

max_batch_size(graph_walk)[source]#
Parameters:

graph_walk (str)

node_resources: dict[str, Any]#
postprocess(request_id, request_info, outputs, inputs=None, **kwargs)[source]#

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:
abstractmethod prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#
Parameters:
Return type:

NodeInputs

preprocess(graph_walk, engine_inputs, inputs)[source]#
Parameters:
Return type:

dict[str, Tensor | Any]

request_state(request_id)[source]#

The request’s state, created on first access.

Parameters:

request_id (str)

Return type:

PerRequestState

request_states: dict[str, PerRequestState]#
unpack_packed_outputs(static_output, request_ids, real_seq_lens, inputs, per_request_info)[source]#

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:
Return type:

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

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

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:
add(key, value)[source]#
Parameters:

key (str)

Return type:

None

add_all(**kwargs)[source]#
Return type:

None

disag_shared_keys: set[str]#
get(key, default=None)[source]#
Parameters:

key (str)

kwargs: dict[str, Any]#
remove(keys)[source]#
Return type:

None

tensors: dict[str, Tensor]#
class mstar.model.submodule_base.StackingMethod(*values)[source]#

Bases: Enum

CAT = 'cat'#
NONE = 'none'#
STACK = 'stack'#
Contents