mstar.engine.engine

Contents

mstar.engine.engine#

Classes

Engine([autocast_dtype, enable_nvtx, ...])

ExecutingBatch(node_name, per_request_info, ...)

SubmoduleManagement(submodule, forward, ...)

class mstar.engine.engine.Engine(autocast_dtype=torch.bfloat16, enable_nvtx=False, enable_profile=False)[source]#

Bases: object

Parameters:

  • enable_nvtx (bool)

  • enable_profile (bool)

add_request(request_id, overrides=None)[source]#
Parameters:
Return type:

None

can_pre_plan(node_name)[source]#

Whether pre-planning this node could stage anything at all.

A node with no pre-planning resource has nothing to stage: its step declaration is empty, so pre_plan_for_batch is guaranteed to bail. Asked BEFORE the plan thread waits on the previous batch’s commit, so a stateless node does not pay that wait to reach a foregone conclusion.

Parameters:

node_name (str)

Return type:

bool

check_ready(node_name, request_id, request_info)[source]#

Whether this node can run the request now.

An offloaded request is brought back first; it stays not-ready until that fits, which is the scheduler’s cue to run something else (and, on OOM, to evict). Then each resource takes in whatever the request published elsewhere — a KV transfer from a prefill worker, say — and reports whether that has landed.

The outcome, not a bool: a resource can report the request unservable (AdmitRuntimeError), and the caller has to fail it rather than scan it again forever.

Parameters:
Return type:

FullAdmitOutcome

check_stop_for_batch(batch, outputs)[source]#

Each rid’s submodule.check_stop, as rid -> loops that should stop.

Reads tensor values, so it belongs on the caller’s slow-postprocess path rather than in exec, where the .item() / .cpu() would block the GPU thread. Per-rid like the other stages: a raise fails that request onto the batch, because letting it escape would abandon the check for the rest of the batch and leave their loops running past their stop condition.

Parameters:
Return type:

dict[str, set[str]]

evictable(node_name)[source]#
Parameters:

node_name (str)

Return type:

bool

exec(batch)[source]#

Run one step: declare → admit → plan → forward → commit.

Captured replay and the eager forward are the same path. Under a lease the batch is padded to the slot’s shape, preprocess output is staged into the static buffers, and the launch is a replay; everything above the launch is identical.

Parameters:

batch (ExecutingBatch)

Return type:

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

exec_and_postprocess(batch)[source]#

The forward and its per-rid tail, which belong to the same step: a walk that captured only its forward finishes in postprocess.

prepare_inputs and the stop check stay outside — the worker has to place those itself.

Parameters:

batch (ExecutingBatch)

Return type:

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

finalize_batch(batch)[source]#
Parameters:

batch (ExecutingBatch)

get_max_batch_size(node_name, graph_walk)[source]#

Most requests this node will take in one step, or None for no cap.

Two sources: what the submodule says it can batch, and the largest batch this walk captured a graph for — going past that would drop the step to eager, so the scheduler splits instead. Splitting is the scheduler’s job, not the engine’s: the pieces then pipeline like any other batch instead of running back to back.

Parameters:

  • node_name (str)

  • graph_walk (str)

Return type:

int | None

is_offloaded(node_name, request_id)[source]#
Parameters:

  • node_name (str)

  • request_id (str)

Return type:

bool

load_model(submodules, specs, parallel_groups, device, transfer_engine_info, kv_cache_type=None)[source]#
Parameters:
offload_priority(node_name, request_id, resource_label)[source]#

How much one named resource wants this request gone.

Only the PRIORITY eviction policy asks; it names the resource to consult, since “most worth reclaiming” means something different per resource.

Parameters:

  • node_name (str)

  • request_id (str)

  • resource_label (str)

Return type:

float

offload_request(node_name, request_id)[source]#

Move the request off-device across this node’s resources.

Returns what was reclaimed, in whatever each resource counts (pages, today); 0 means nothing moved and the caller should pick another victim.

Parameters:

  • node_name (str)

  • request_id (str)

Return type:

int

postprocess_batch(batch, outputs)[source]#

Per-rid submodule.postprocess, e.g. the non-capturable tail of a walk whose graph covered only the forward.

Per-rid like prepare_inputs: a raise fails that request and leaves the rest of the batch to route normally.

Parameters:
Return type:

None

pre_plan_for_batch(batch)[source]#

Admit and plan the pre-planning resources a step ahead.

Worth doing only under a lease: the reserved slot is not the one the in-flight replay reads, so this can write its plan buffers on the plan stream while that replay runs. exec then re-drives the full sweep, where each pre-planned resource promotes what was staged here, and the replay waits on the event recorded here.

An unprepared batch declares its step over the capture config’s inputs for the leased bucket, which is why the lease had to come from a batched capture. That step is not cached: exec re-declares it over the real inputs.

TODO: pre-planning a packed capture (and chunked prefill) needs the real token counts here, i.e. prepare_inputs run ahead of the forward. That needs every submodule’s prepare_inputs to be async-safe — no .item() on a tensor the in-flight step is still producing — which some are not (e.g. qwen3_tts). ExecutingBatch.outputs_ready is the hook for it.

Returns False when nothing was planned ahead, in which case exec plans inline.

Parameters:

batch (ExecutingBatch)

Return type:

bool

prepare_inputs(batch)[source]#

Per-rid submodule.prepare_inputs, onto batch.inputs.

Per-rid, so a raise is attributable to one request: record it and take that rid out rather than losing the batch. A rid the submodule declines (returns None) leaves the same way, without being an error.

Parameters:

batch (ExecutingBatch)

Return type:

None

preplan_is_stale(batch)[source]#

Whether the staged plan no longer describes this batch.

Parameters:

batch (ExecutingBatch)

Return type:

bool

reclaimable(node_name, request_id, affected_resources=None)[source]#

What this node’s resources could reclaim from the request. 0 means offloading it would free nothing, so it is not a candidate.

Parameters:

  • node_name (str)

  • request_id (str)

  • affected_resources (set[str] | None)

Return type:

int

reload_request(node_name, request_id)[source]#

Bring it back. False when any resource can’t fit it yet, in which case the request stays offloaded and the caller retries later.

Parameters:

  • node_name (str)

  • request_id (str)

Return type:

bool

remove_request(request_id)[source]#
Parameters:

request_id (str)

Return type:

None

reserve_replay_slot(batch)[source]#

Lease the slot this batch will replay on, before it is dispatched.

Reserving up front is what puts pre-plan(N+1) and replay(N) on different slots. No captured graph for the batch’s shape leaves the lease unset, i.e. the step runs eager.

A batch that hasn’t been through prepare_inputs has no token count yet, so only a batched capture can serve it — see CudaGraphRunner.select_batched_bucket.

Parameters:

batch (ExecutingBatch)

Return type:

SlotLease | None

reset_pre_plan_for_batch(batch=None)[source]#

Drop planned-ahead state, e.g. when its batch never dispatched.

The resources hold the plan itself. The batch keeps whatever prepare_inputs produced — only the plan is redone — but the leased slot’s padding rows are returned so their pages aren’t left attributed to a step that never ran.

Parameters:

batch (ExecutingBatch | None)

Return type:

None

shutdown()[source]#
warmup()[source]#
Return type:

None

class mstar.engine.engine.ExecutingBatch(node_name: str, per_request_info: Mapping[str, mstar.conductor.request_info.CurrentForwardPassInfo], step_context: mstar.engine.resources.step.StepContext, running_batched: bool = False, cg_key_info: typing.Any | None = None, preplan_event: torch.cuda.streams.Event | None = None, preplanned_rids: tuple[str, ...] | None = None, step: mstar.engine.resources.step.SubmoduleStep | None = None, per_request_input_tensors: Mapping[str, dict[str, list[torch.Tensor]]] = <factory>, final_stream_rids: set[str] = <factory>, inputs: list[mstar.model.submodule_base.NodeInputs] | None = None, skipped_rids: set[str] = <factory>, failed_requests: dict[str, str] = <factory>, admit_error: mstar.engine.resources.step.AdmitFailedReason | None = None, failed_resource: str | None = None, outputs: dict[str, dict[str, list[torch.Tensor]]] = <factory>, outputs_ready: threading.Event = <factory>, commit_done: threading.Event = <factory>, completion_event: torch.cuda.streams.Event | None = None, launch_started_event: threading.Event | None = None, exec_timings: mstar.profile.worker.ExecTimings = <factory>)[source]#

Bases: object

Parameters:
admit_error: AdmitFailedReason | None = None#
cg_key_info: Any | None = None#
commit_done: Event#
completion_event: Event | None = None#
drop_rids(rids)[source]#

Take rids out of the step, leaving the rest of the batch intact.

Parameters:

rids (set[str])

exec_timings: ExecTimings#
failed_requests: dict[str, str]#
failed_resource: str | None = None#
final_stream_rids: set[str]#
property graph_walk: str#
inputs: list[NodeInputs] | None = None#
launch_started_event: Event | None = None#
lease_slot(slot_lease)[source]#
Parameters:

slot_lease (SlotLease)

node_name: str#
outputs: dict[str, dict[str, list[Tensor]]]#
outputs_ready: Event#
per_request_info: Mapping[str, CurrentForwardPassInfo]#
per_request_input_tensors: Mapping[str, dict[str, list[Tensor]]]#
preplan_event: Event | None = None#
preplanned_rids: tuple[str, ...] | None = None#
register_admit_error(reason, failed_resource=None)[source]#
Parameters:
register_failure(rid, error)[source]#
Parameters:
register_prepare_batch(inputs)[source]#
Parameters:

inputs (list[NodeInputs])

release_waiters()[source]#

Let anything waiting on this step proceed.

Called on every exit from exec, so a step that raised before publishing outputs or committing doesn’t strand the thread preparing the next one.

property request_ids#
running_batched: bool = False#
skipped_rids: set[str]#
step: SubmoduleStep | None = None#
step_context: StepContext#
class mstar.engine.engine.SubmoduleManagement(submodule: mstar.model.submodule_base.NodeSubmodule, forward: Callable, forward_batched: Callable, joint_comm_group: mstar.distributed.communication.JointGroups, resources: dict[str, mstar.engine.resources.base.Resource], cuda_graph_runner: mstar.engine.cuda_graph_runner.CudaGraphRunner | None = None, piecewise_runners: dict[str, mstar.engine.cuda_graph_runner.PiecewiseCudaGraphRunner] = <factory>)[source]#

Bases: object

Parameters:
cuda_graph_runner: CudaGraphRunner | None = None#
forward: Callable#
forward_batched: Callable#
joint_comm_group: JointGroups#
piecewise_runners: dict[str, PiecewiseCudaGraphRunner]#
resources: dict[str, Resource]#
submodule: NodeSubmodule#
Contents