mstar.engine.engine

Contents

mstar.engine.engine#

Functions

checkpoint_identity(path)

Name the weights a cached page was produced under, without reading them.

Classes

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

ExecutingBatch(node_name, per_request_info, ...)

ForwardPassInfoWrapper(info, metadata)

Stop-gap solution: the dynamic loop iters in CurrentForwardPassInfo may be used in two places at once: in _postprocess_batch (for batch N), and in the forward pass for batch N+1, which require different iterations.

SubmoduleManagement(submodule, forward, ...)

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

Bases: object

Parameters:

  • graph_runtime (GraphRuntime)

  • enable_nvtx (bool)

  • enable_profile (bool)

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

None

apply_resident_delta(node_name, delta)[source]#

Match the leader’s resident set on this node. Mutates OffloadDelta by popping off the steps that have already landed. False means not all has landed and the caller must retry before trying the next step.

Parameters:
Return type:

bool

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

capture_group(node_name, graph_walk, request_id, request_info)[source]#

The capture bucket this one request would replay, by cg_key_info; None when the walk has no capture or the submodule does not split by it.

A batch spanning two groups matches no capture and runs eager, so the scheduler batches each group on its own when the submodule opts in (split_batches_by_capture_key).

Parameters:
Return type:

Any | None

check_ready(node_name, request_id, request_info, allow_reload=True)[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, host_rows=None)[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.

With host_rows, check_stop_batched checks the batch in one call; if it raises, the per-request path runs so a failure lands on its own request.

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:

BatchedModelOutput

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:

BatchedModelOutput

extend_prefix_chains(batch, outputs)[source]#

Key what this step generated, from the stop check’s host copy.

Parameters:
Return type:

None

finalize_batch(batch, publish_request_ids=None)[source]#
Parameters:
Return type:

dict[str, dict[str, PublishedInfo]]

finalize_stopped_requests(batch, request_ids)[source]#

Publish state that must include the last iteration of a stopped loop.

Parameters:
Return type:

None

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, model=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 (int)

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

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

The moves on this node since the last call, ready for the wire.

Verbatim. Every move in it is one this rank executed, so a rank in the same state can always apply it — which is what makes a refused replay mean the states differ, and nothing else.

Parameters:

node_name (str)

Return type:

OffloadDelta

warmup()[source]#
Return type:

None

class mstar.engine.engine.ExecutingBatch(node_name: str, per_request_info: Mapping[int, mstar.conductor.request_info.CurrentForwardPassInfo], step_context: mstar.engine.resources.step.StepContext, running_batched: bool = False, slot: int | None = None, cg_key_info: typing.Any | None = None, preplan_event: torch.cuda.streams.Event | None = None, preplanned_rids: tuple[int, ...] | 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>, per_request_input_metadata: Mapping[str, mstar.model.submodule_base.InputMetadata] = <factory>, stream_partition_done_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: mstar.model.submodule_base.BatchedModelOutput = <factory>, resource_publish_info: dict[str, dict[str, mstar.engine.resources.base.PublishedInfo]] = <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: BatchedModelOutput#
outputs_ready: Event#
per_request_info: Mapping[int, CurrentForwardPassInfo]#
per_request_input_metadata: Mapping[str, InputMetadata]#
per_request_input_tensors: Mapping[str, dict[str, list[Tensor]]]#
preplan_event: Event | None = None#
preplanned_rids: tuple[int, ...] | 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#
resource_publish_info: dict[str, dict[str, PublishedInfo]]#
running_batched: bool = False#
set_slot(slot)[source]#
Parameters:

slot (int)

skipped_rids: set[str]#
slot: int | None = None#
step: SubmoduleStep | None = None#
step_context: StepContext#
stream_partition_done_rids: set[str]#
class mstar.engine.engine.ForwardPassInfoWrapper(info, metadata)[source]#

Bases: object

Stop-gap solution: the dynamic loop iters in CurrentForwardPassInfo may be used in two places at once: in _postprocess_batch (for batch N), and in the forward pass for batch N+1, which require different iterations. The InputMetadata is built per executing batch, so its value is guaranteed to be up-to-date.

TODO: the long-term solution is to maybe move the dynamic loop iters off of CurrentForwardPassInfo entirely.

Parameters:
info: CurrentForwardPassInfo#
metadata: InputMetadata#
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#
lease_slot()[source]#
Return type:

int

property needs_slot_fence: bool#

Whether reusing a slot has to wait out the GPU work that last held it — true exactly when a resource stages into a reused host buffer.

property next_slot: int#

an int read is atomic under the GIL, and the only other leaser is the plan thread, which is gated on the previous batch’s commit_done.

Type:

The slot the next lease hands out. Read unlocked

property num_slots: int#
piecewise_runners: dict[str, PiecewiseCudaGraphRunner]#
resources: dict[str, Resource]#
set_piecewise_slot(slot)[source]#

Point the region runners at this batch’s slot, on the GPU thread just before its forward. Set here rather than at lease time: leasing runs ahead of the forward (pre-plan reserves N+1’s slot while N is in flight), so a lease-time write would race the replay reading _current_slot.

Parameters:

slot (int)

Return type:

None

submodule: NodeSubmodule#
mstar.engine.engine.checkpoint_identity(path)[source]#

Name the weights a cached page was produced under, without reading them.

A checkpoint with no safetensors index falls back to its shard names and sizes, which catches a replaced shard but not one edited in place at the same length.

Parameters:

path (str | Path)

Return type:

bytes

Contents