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

BatchedModelOutput([per_rid_outputs, ...])

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

BatchedStopCheck(stops, withheld)

StopCheck for check_stop_batched: rid -> loops to stop and rid -> withheld outputs, rids with none left out of each.

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)

StopCheck([stops, withheld])

check_stop's answer when the step also has outputs that are not data.
class mstar.model.submodule_base.ARNodeInputs(tensor_inputs=<factory>, kwargs=<factory>, resource_step_info=None, input_seq_len=0, graph_walk=None, 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

cuts_tensor_inputs: bool = False#
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]

split_inputs(graph_walk, fwd_info, inputs, start, end)[source]#

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:
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, row_views=None, host_rows=None, rows_only=False, loopback_rids=frozenset({}), output_spans=None, span_request_ids=None, host_spans=None)[source]#

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:
check_stop_buffers: dict[str, Tensor | dict[str, list[Tensor]]] | None = None#
clone_check_stop_buffers(reuse=None)[source]#

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

reuse holds clones already taken of the row outputs (see clone_row_outputs); a stop buffer that is the very same tensor as a row output takes that clone instead of a second copy, and a tensor under two names is cloned once. The row clone is narrowed to the real rows, which is every row the stop check reads.

Parameters:

reuse (dict[str, Tensor] | None)

clone_row_outputs(num_rows)[source]#

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

output (BatchedModelOutput | dict[str, Any])

Return type:

BatchedModelOutput

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

key (str)

get_check_stop_input()[source]#
host_rows: HostRows | None = None#
host_spans: dict[str, Tensor] | None = None#
loopback_rids: frozenset = frozenset({})#
output_spans: dict[str, Tensor] | None = None#
packed_outputs: dict[str, Tensor]#
per_rid_outputs: dict[str, dict[str, list[Tensor]]]#
pop(request_id, default=None)[source]#

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#
row_views: dict[str, tuple[Tensor, ...]] | None = None#
rows_only: bool = False#
span_request_ids: tuple | None = None#
update(other)[source]#
Parameters:

other (BatchedModelOutput | dict[str, Any])

class mstar.model.submodule_base.BatchedStopCheck(stops, withheld)[source]#

Bases: NamedTuple

StopCheck for check_stop_batched: rid -> loops to stop and rid -> withheld outputs, rids with none left out of each.

Parameters:
stops: dict[str, set[str]]#

Alias for field number 0

withheld: dict[str, tuple[str, ...]]#

Alias for field number 1

class mstar.model.submodule_base.HostRows(request_ids, buffers)[source]#

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:
buffers: dict[str, Tensor]#

Alias for field number 1

first_values(name, request_ids)[source]#

(rid, value) for each of request_ids that has a row, where value is the first element of its row of buffer name, off one tolist; None when the step staged no such buffer. A request with no row is skipped, as the per-request stop path skips a request with no output.

Parameters:
request_ids: tuple#

Alias for field number 0

class mstar.model.submodule_base.InputMetadata(stream_chunks=mappingproxy({}), dynamic_loop_iter_counts=mappingproxy({}))[source]#

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:
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)[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, 'PiecewiseAcceleratorGraphRunner']"=<factory>, per_request_states: "'Mapping[str, PerRequestState] | None'"=None, step: 'SubmoduleStep | None' = None, captured: 'bool' = False, synthetic: 'bool' = False, real_seq_len: 'int | None' = None, per_request_input_metadata: 'Mapping[str, InputMetadata]'=<factory>)[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_input_metadata: Mapping[str, InputMetadata]#
per_request_states: Mapping[str, PerRequestState] | None = None#
piecewise_runners: dict[str, 'PiecewiseAcceleratorGraphRunner']#
real_seq_len: int | None = None#
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#
synthetic: bool = False#
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, graph_walk: 'str | None' = None)[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.

graph_walk: str | None = None#
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.

batch_postprocess(outputs, request_infos, inputs, piecewise_runners=None)[source]#

Once per step, for work that batches across the step’s requests (e.g. watermarking the audio chunks postprocess trimmed). Works on outputs in place: its rows on a rows-only step, which has no per-request outputs, and otherwise the per-request outputs after every request’s postprocess, for the requests whose postprocess succeeded (request_infos’ keys, in the order of inputs); a raise then fails every one of them. Optional; only called when overridden.

piecewise_runners are this node’s captured regions (get_piecewise_accelerator_graph_configs), as engine_inputs hands them to the forward, so this work can replay a capture too.

Parameters:
Return type:

None

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_accelerator_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, per_request_input_metadata=None, **kwargs)[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.

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

A StopCheck also names outputs of this step that are not routed (StopCheck.withheld), decided off the same read; it may withhold without stopping.

Default: no stops.

Parameters:
Return type:

set[str] | StopCheck

check_stop_batched(request_ids, request_infos, host_rows, step_walk=None)[source]#

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

Optional. Returns rid -> loops to stop (rids with none left out), or a BatchedStopCheck for a submodule whose check_stop withholds, and must agree with check_stop for every rid. None, or a raise, falls back to the per-request path. step_walk is the walk every row ran, or None when the step mixes walks (each row’s is on its info).

Parameters:
Return type:

dict[str, set[str]] | BatchedStopCheck | None

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)

cudnn_allow_tf32: bool | None = None#
declare_step(graph_walk, request_ids, inputs, slot_lease=None, piecewise_leases=None, per_request_info=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.

per_request_info holds the real rows’ forward-pass info (capture passes dummy metadata); padding rows are absent, so read it with .get.

Parameters:
Return type:

SubmoduleStep | None

device_loopback_signals(graph_walk)[source]#

Loop-back input signals of this walk whose per-row value the node keeps on the device itself (a sampler’s slot master), so the worker routes them with no tensor: the next step reads the value in preprocess and every fallback batch arrives with the signal empty. prepare_inputs is told per row (has_loopback) whether its last step here kept the value on the device; a row without it (e.g. its previous step ran on another rank) brings the tensor. On a walk whose output signals all stay on the device the engine builds no per-request output dicts and does not call postprocess: the stop check, the inline emit and the next step’s speculation read the step’s rows instead. Empty by default; MSTAR_DEVICE_LOOPBACK=0 keeps every node on the tensor path.

Parameters:

graph_walk (str)

Return type:

frozenset[str]

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]]

float32_matmul_precision: str | None = None#
abstractmethod forward(graph_walk, engine_inputs, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable. A BatchedModelOutput keyed by the request id, as forward_batched returns, can carry output_spans.

Parameters:
Return type:

dict[str, list[Tensor]] | BatchedModelOutput

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

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

Parameters:
Return type:

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

get_accelerator_graph_configs(device, tp_world_size=1)[source]#
Parameters:
Return type:

list[AcceleratorGraphConfig]

get_autocast_dtype()[source]#

This node’s dtype in place of the model’s: the engine manager casts the node’s params to it, and every forward, warm-up and capture runs under autocast in it. Read once at load (resolve_node_dtype).

Default: None (inherit the model’s dtype). Ignored when disable_autocast is set; to run a node without autocast, in the dtype it was built in, set that instead.

Return type:

dtype | None

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

Return the piecewise CUDA graph configs this submodule opts into.

autocast_dtype is this node’s resolved dtype, float32 for a node with disable_autocast — 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, PiecewiseAcceleratorGraphConfig]

inline_client_signals(graph_walk)[source]#

Output signals whose client-facing value is one scalar per row that the stop check already copies to the host, as signal name -> the name of the row-addressed stop buffer holding it. The worker then sends those values inline, one frame per step, instead of one tensor per request. Only for a walk whose emit edge is not persisted.

Parameters:

graph_walk (str)

Return type:

dict[str, str]

max_batch_size(graph_walk)[source]#
Parameters:

graph_walk (str)

max_rows_in_combined(graph_walk)[source]#

How many rows of this real walk a step mixing it with another walk of its combined walk may hold; None is uncapped. For a backend whose mixed step runs the walk’s rows on a path sized to a few of them.

Parameters:

graph_walk (str)

Return type:

int | None

node_resources: dict[str, Any]#
on_request_added(request_id, request_info)[source]#

The request has reached this worker; none of its steps has run yet.

Runs on the worker’s main thread, so host work worth starting early belongs on a thread of the submodule’s own, and nothing here may touch the device. State goes in request_state(), which the removal of the request frees. Default: nothing.

Parameters:
Return type:

None

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]

reads_device_loopback: bool = False#
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]#
spanned_outputs(graph_walk)[source]#

Outputs of this walk that may come with a valid span per row, in BatchedModelOutput.output_spans under the output’s name. The worker cuts each row to its span along the last dim once the step’s completion event has fired, before the stop check, transport or a streaming edge reads it, so a forward (and its captured graphs) can keep a fixed width and still emit a length known only on the device, without a host read. The cut is a view, and a row whose span covers its whole output is left as it is.

The engine-side postprocess and batch_postprocess see the uncut output. A spanned output must not feed a node this worker speculates into, which would read it before the cut; startup refuses that unless one of the two nodes sets enable_async_scheduling=False. Not for a signal kept on the device or sent inline, which have no per-request tensor to cut. Empty by default.

Parameters:

graph_walk (str)

Return type:

frozenset[str]

split_batches_by_capture_key(graph_walk)[source]#

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

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

NodeInputs

step_is_input_free(graph_walk)[source]#

Whether declare_step for this walk reads nothing from the inputs but their spans, and every real row has the span of the template row. Then the step a pre-plan declared over the bucket’s template rows is the step the forward would declare, and the engine runs it instead of declaring and admitting a second time.

Parameters:

graph_walk (str)

Return type:

bool

uniform_row_inputs(graph_walk, has_loopback=False)[source]#

The inputs every row of this walk gets, when prepare_inputs would return an equal object for each regardless of the request, its routed signals and its metadata (a decode walk that reads its token off the device). The engine then prepares the batch with one call and shares the object across the step’s rows, so it must not be mutated per row. has_loopback: every row’s last step on this node kept its loop-back signals on the device (device_loopback_signals). None (the default) prepares row by row.

Parameters:

  • graph_walk (str)

  • has_loopback (bool)

Return type:

NodeInputs | None

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]]]

warmup_compiles()[source]#

Compile what a forward would otherwise compile on its first request, outside the accelerator graph runner. Called once, at the end of the engine’s warmup, after capture, under no_grad and the node’s fp32 precision; the node’s resources are bound but nothing is admitted or planned, so it must not read or change resource state. On a lockstep node every rank calls it in the same order, so it may run the collectives its forward runs. Default: nothing.

Return type:

None

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'#
class mstar.model.submodule_base.StopCheck(stops=frozenset({}), withheld=())[source]#

Bases: NamedTuple

check_stop’s answer when the step also has outputs that are not data.

withheld names outputs of this step the worker leaves unrouted for the request, e.g. a frame streamed beside a stop token: a streaming edge then gets no item for it, so its consumer needs no device read to find and drop it. A plain set return is the same as StopCheck(stops).

Parameters:
stops: set[str] | frozenset[str]#

Alias for field number 0

withheld: tuple[str, ...]#

Alias for field number 1

Contents