mstar.engine.accelerator_graph_runner

Contents

mstar.engine.accelerator_graph_runner#

Functions

agree_across_ranks(comm_group, values, ...)

Reduce each of values over every rank of the joint group with reduce.

autocast_scope(dtype, device_type)

A forward's autocast scope; None (disable_autocast) runs the submodule in its own dtype and shuts out any ambient autocast.

blocked_capture_resources(step, resources)

Return resource keys that prevent capture of the declared step.

capture_into_graph(run, pool, device, ...[, ...])

Capture run into a new CUDA graph allocated from pool.

capture_with_static_outputs(run, ...[, ...])

Capture run so its outputs land in buffers allocated OUTSIDE pool.

compile_kwargs(config_mode)

torch.compile's arguments for a capture config whose compile_mode is config_mode, which the config checked when it was built.

dummy_metadata(rids, graph_walk)

Stand-in request info for padding rows, which have no real request.

dummy_rid_handle(name)

dummy_rid_name(handle)

The name a padding-row handle was minted from.

fail_if_graphs_required(missing)

MSTAR_REQUIRE_CUDA_GRAPHS=1 turns a dropped bucket into a startup failure, for deployments that would rather not come up than serve the eager path at 10-20x the latency.

plan_captures(pools, free[, reusable])

The cheapest buckets across every pool on a device, added while the largest eager step still fits in free beside them.

Classes

AcceleratorGraphBucket(config, config_idx[, ...])

The captured slots for one (walk, cg_key_info, bs, num_tokens).

AcceleratorGraphRunner(submodule_name, ...)

AcceleratorGraphSlot(graph, static_inputs, ...)

One captured graph + the buffers its replay reads and writes.

CaptureCost(peak, graph, kept, slots[, ...])

What one bucket costs on the device: its warm-up, and a capture thrown away, its own or its walk's largest.

CapturePlan(buckets, predicted, outside, floor)

Which buckets each graph pool on a device captures, chosen before any is.

DummyRowPool(prefix, step_runner, resources)

The padding rows captured replays pad onto.

ExecutableMemory()

The executable memory a walk's graphs take per node, outside any pool, measured on copies of its first throwaway graph: executables instantiated from it through libcuda.

GraphMemory(pool, outside)

What a captured graph holds: in its pool, and outside the allocator.

PiecewiseAcceleratorGraphRunner(label, ...)

Captures one inner callable of a submodule's forward as a CUDA graph.

PiecewiseGraphData(graph, static_inputs, ...)

One captured region, for one (bs, total_tokens) bucket.

PiecewiseGraphKey(bs, seq_len, slot)

PiecewiseOutput(outputs, real_len)

Dict-like view over a captured region's output buffers.

WarmedRegion(run, static_inputs, dummy_rids, ...)

A region's shape planned and run eagerly: what its capture reads.

WarmedSpec(run, static_inputs, ...)

A spec prepared and run eagerly: what its capture reads.
class mstar.engine.accelerator_graph_runner.AcceleratorGraphBucket(config, config_idx, slots=<factory>)[source]#

Bases: object

The captured slots for one (walk, cg_key_info, bs, num_tokens).

Parameters:
config: AcceleratorGraphConfig#
config_idx: int#
slots: list[AcceleratorGraphSlot]#
class mstar.engine.accelerator_graph_runner.AcceleratorGraphRunner(submodule_name, submodule, resources, step_runner, device, autocast_dtype, joint_comm_group, num_slots, enable_nvtx=False, device_capability=None)[source]#

Bases: object

Parameters:
CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32, 64]#
NUM_WARMUP = 2#
property any_graphs#
can_run(graph_walk, bs, num_tokens, cg_key_info=None)[source]#
Parameters:

  • graph_walk (str)

  • bs (int)

  • num_tokens (int)

  • cg_key_info (Any | None)

Return type:

bool

captures_walk(graph_walk)[source]#

Whether any capture replays this walk.

Parameters:

graph_walk (str)

Return type:

bool

config_for(lease)[source]#
Parameters:

lease (SlotLease)

Return type:

AcceleratorGraphConfig

declare_inputs_for(lease)[source]#

Template rows for a declare-only call, cached per bucket.

A bucket’s shape is fixed, so these rows are constant; building them fresh cost ~130us of clones per pre-plan at bs=16, plus that many D2D copies. Only for callers that hand the rows to declare_step and nothing else — the list is shared, so a caller that mutates it or passes it to preprocess corrupts every later step on this bucket.

Parameters:

lease (SlotLease)

Return type:

list[NodeInputs]

dropped_buckets: list[BucketKey]#
is_packed(lease)[source]#

Whether lease replays a packed capture.

Parameters:

lease (SlotLease)

Return type:

bool

lease_slot(graph_walk, bs, slot, num_tokens=None, cg_key_info=None)[source]#

Pick the bucket and double-buffer slot an upcoming step replays on.

num_tokens=None means the batch’s inputs aren’t built yet, so the bucket comes from the batched-capture search instead.

slot=None advances the bucket’s counter so the next lease lands on the other slot; a caller that already reserved one (pre-plan) passes it back so both submissions target the same slot.

Parameters:

  • graph_walk (str)

  • bs (int)

  • slot (int)

  • num_tokens (int | None)

  • cg_key_info (Any | None)

Return type:

SlotLease | None

max_batch_size_for(graph_walk)[source]#

Largest batch this walk was captured for, or None for no cap.

A config that opts out of capping captures only an acceleration subset of batch sizes — larger batches run eager rather than being split — so it does not constrain the batch. With every config for the walk opting out, nothing here caps it.

Parameters:

graph_walk (str)

Return type:

int | None

pad_inputs(lease, inputs)[source]#

Real inputs plus the rows that bring the batch to capture batch size. Set the length to zero so no new pages are allocated for dummy requests.

Parameters:
Return type:

list[NodeInputs]

plan_stream()[source]#

Dedicated stream for pre-planning.

Pre-plan must not submit onto the default stream: whether its memcpys land before or after the GPU thread records the previous batch’s completion event is timing-dependent, and landing after delays that event past pre-plan’s own kernels. Its own stream keeps the two independent; a plan-done event gates the replay that reads the buffers.

Return type:

Any | None

prepare_for_capture()[source]#

Claim the static buffers this runner’s captures will read.

Split from the capture so every runner claims first: nodes share resources, and a build driven by a later node would move buffers an earlier node’s graphs already baked in.

Return type:

list[CGSlotSpec]

release(lease, real_bs)[source]#

Return the padding rows to their at-rest state after a step.

Padding rows (is_padding_row) run against SINK_PAGE and the recurrent pool’s sink slot, so nothing stays resident: padding must never compete with real requests for storage.

Parameters:
Return type:

None

run_forward(lease, preprocessed, plan_done_event=None, launch_started_event=None)[source]#

Stage this step’s inputs into the slot and replay it.

plan_done_event is the pre-plan’s event: its plan wrote this slot’s buffers on another stream, so the replay has to wait for those writes before it reads them.

launch_started_event releases the submitting thread. It is set HERE rather than before the forward: staging copies each preprocessed tensor into its static buffer and holds the GIL while doing so, so releasing earlier hands the main thread a GIL the GPU thread still needs. Only graph.replay() below drops it in C++.

Parameters:
Return type:

dict

select_batched_bucket(graph_walk, bs, cg_key_info=None)[source]#

A bucket for a batch whose inputs aren’t built yet.

Only a batched capture can answer this: its token count is a property of the config (bs rows of a fixed per-request length), so the bucket follows from the batch size alone. A packed capture’s token count is a property of the requests, so it has to wait for prepare_inputs and go through select_bucket.

This is what lets pre-plan run before inputs exist. Lifting the restriction means making prepare_inputs safe to run ahead of the forward — see Engine.pre_plan_for_batch.

Parameters:

  • graph_walk (str)

  • bs (int)

  • cg_key_info (Any | None)

Return type:

BucketKey | None

select_bucket(graph_walk, bs, num_tokens, cg_key_info=None)[source]#

Tightest captured bucket that fits this batch, or None for eager.

A walk may have several captures (e.g. one per image resolution, each a fixed shape with its own token count), so every matching bucket is considered rather than the first config declared. num_tokens counts the batch’s input tokens; a bucket is keyed by the tokens its plan carries, so the request is scaled by that config’s multiplier.

Parameters:

  • graph_walk (str)

  • bs (int)

  • num_tokens (int)

  • cg_key_info (Any | None)

Return type:

BucketKey | None

size_captures()[source]#

Warm every spec up, and capture each walk’s largest buckets once, thrown away, before any runner captures for real.

Every bucket is warmed rather than the first in capture order: that order leads with batch size, so it can open on a decode bucket while a one-row prefill of thousands of tokens peaks far higher. And every slot: a slot’s first plan builds state of its own (FlashInfer’s graph wrappers and their workspace, 750 MiB for the Talker’s second slot), which has to be spent before free memory is read.

Return type:

dict[BucketKey, CaptureCost]

size_lent_captures(wanted)[source]#

Capture every bucket a lender sized, thrown away, if wanted here or on another rank: a plan that leaves graphs out needs what each takes, not what its walk’s largest does.

Every rank asks, so the ranks of a group size again together or not at all, and plan on the same kind of cost. With no forward and no barrier: sizing ran the forwards, and a capture runs no collective.

Parameters:

wanted (bool)

Return type:

dict[BucketKey, CaptureCost]

slot_for(lease)[source]#
Parameters:

lease (SlotLease)

Return type:

AcceleratorGraphSlot

step_ids(lease, request_ids)[source]#

The padded addressing for one step: real ids first, then the slot’s padding ids. Plans, commits, and advances address real request state by its own id; only the padding rows run against the slot’s own state.

Parameters:
Return type:

list[int]

step_metadata(lease, request_ids, per_request_info)[source]#
Parameters:
Return type:

dict[int, CurrentForwardPassInfo]

warmup_and_capture(planned)[source]#

Capture the planned buckets, of every config and batch size. With no plan, every bucket, each warmed up first as sizing would have.

Parameters:

planned (Collection[BucketKey] | None)

class mstar.engine.accelerator_graph_runner.AcceleratorGraphSlot(graph, static_inputs, static_input_keys, static_outputs, dummy_rids, dummy_metadata, config_idx)[source]#

Bases: object

One captured graph + the buffers its replay reads and writes.

Double-buffer: a bucket holds one of these per slot. Replay alternates between slots so plan(N+1) on the inactive slot’s resources can run concurrently with replay(N) on the active slot — so a node with nothing to pre-plan keeps a single slot.

Parameters:
config_idx: int#
dummy_metadata: dict[int, CurrentForwardPassInfo]#
dummy_rids: list[int]#
graph: CapturedGraph#
static_input_keys: tuple[str, ...]#
static_inputs: dict[str, Any]#
static_outputs: BatchedModelOutput#
class mstar.engine.accelerator_graph_runner.CaptureCost(peak, graph, kept, slots, outside=0, executable=None)[source]#

Bases: NamedTuple

What one bucket costs on the device: its warm-up, and a capture thrown away, its own or its walk’s largest.

Parameters:
executable: ExecutableMemory | None#

Alias for field number 5

graph: int | None#

Alias for field number 1

kept: int#

Alias for field number 2

outside: int#

Alias for field number 4

peak: int#

Alias for field number 0

slots: int#

Alias for field number 3

class mstar.engine.accelerator_graph_runner.CapturePlan(buckets, predicted, outside, floor)[source]#

Bases: NamedTuple

Which buckets each graph pool on a device captures, chosen before any is.

Parameters:
buckets: dict[str, set]#

Alias for field number 0

floor: int#

Alias for field number 3

outside: dict[str, int]#

Alias for field number 2

predicted: dict[str, int]#

Alias for field number 1

class mstar.engine.accelerator_graph_runner.DummyRowPool(prefix, step_runner, resources)[source]#

Bases: object

The padding rows captured replays pad onto.

Ingested once and kept for the runner’s lifetime: a re-ingest would hand the resources a fresh stream and orphan the pages the last capture left resident, and keeping them resident is what lets a step’s padding tail allocate nothing.

Parameters:
ensure(key, bs)[source]#

bs rows for key, ingesting any this pool hasn’t opened yet.

Parameters:
Return type:

list[int]

handles(key, bs)[source]#

bs padding-row handles for key. Negative and stable per (prefix, key, index), so they never collide with a real request’s handle nor with another pool’s rows.

Parameters:
Return type:

list[int]

release_all()[source]#

Hand back every padding row’s storage once capture is done: a replay pads with zero-length rows, so it is capture-time residue.

Return type:

None

reset(rids, free=False)[source]#
Parameters:
Return type:

None

class mstar.engine.accelerator_graph_runner.ExecutableMemory[source]#

Bases: object

The executable memory a walk’s graphs take per node, outside any pool, measured on copies of its first throwaway graph: executables instantiated from it through libcuda. The cost per node changes with a graph’s size, even for the same kernels: a chain of adds takes 3.1 KiB a node at 300 nodes and 2.4 at 2222. A walk’s graphs are about the same size, so the first graph’s cost per node holds for all of them.

The copies are held until the runner captures. The driver keeps a destroyed graph’s executable memory to reuse, and maps new memory beside it, so after any destroy, free memory drops by less than a copy takes.

measure(graph, nodes)[source]#

Set per_node and held from copies of graph, held until release.

The count starts after the first copy that takes free memory, since copies until then fill pieces the driver mapped earlier. It adds one 2 MiB piece, since the first counted copy may start partway into a piece. Free memory is the device’s, so another process allocating or freeing moves it: a count is kept only if free memory never rose and no copy took more than two pieces above the mean. If no count is kept, per_node is the largest rate a count gave. None if no copy could be counted.

Parameters:

nodes (int)

Return type:

None

per_node: float | None#
release()[source]#
Return type:

None

class mstar.engine.accelerator_graph_runner.GraphMemory(pool, outside)[source]#

Bases: NamedTuple

What a captured graph holds: in its pool, and outside the allocator.

Parameters:
outside: int#

Alias for field number 1

pool: int#

Alias for field number 0

class mstar.engine.accelerator_graph_runner.PiecewiseAcceleratorGraphRunner(label, config, resources, step_runner, device, autocast_dtype, num_slots, joint_comm_group=None, node_name=None, memory_pool=None, device_capability=None)[source]#

Bases: object

Captures one inner callable of a submodule’s forward as a CUDA graph.

Where CudaGraphRunner replays a whole forward_batched under engine control, this captures a SUB-REGION — a transformer block loop, say — while the surrounding preamble stays eager and the submodule invokes run itself. The config supplies the callable, its static buffers, and the region’s own step declaration; the runner drives that step through the same admit → plan → commit cycle the engine uses, so the region’s attention plan lands in the resources’ per-(bucket, slot, label) buffers and the captured graph reads them at fixed addresses.

Parameters:
CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32, 64]#
NUM_WARMUP = 2#
property any_graphs: bool#
can_lease(inputs)[source]#

The config’s can_lease for this batch; True when it sets none.

Parameters:

inputs (list[NodeInputs])

Return type:

bool

can_run(batch_size, total_tokens=None)[source]#
Parameters:

  • batch_size (int)

  • total_tokens (int | None)

Return type:

bool

dropped_shapes: list[tuple[int, int]]#
property lease_before_step: bool#

Whether this region takes its slot ahead of the declaration.

lease_slot(batch_size, total_tokens=None)[source]#

The bucket this batch would replay on, asked before the forward.

Same question run settles per call, answered early so the outer declare_step can plan the region’s resources against it.

Parameters:

  • batch_size (int)

  • total_tokens (int | None)

Return type:

SlotLease | None

prepare_for_capture()[source]#

Claim this region’s static buffers; see CudaGraphRunner.

Return type:

list

run(static_inputs, request_ids=None, seq_lens=None, real_bs=None)[source]#

Replay the captured region for these real inputs.

Copies each real input into the runner-owned buffer of the same name, declares and plans the region’s step over the padded batch, replays, then commits and returns the padded buffers behind a real-length view.

Only the static buffers carry data into a replay: the region’s Python ran once, at capture, so whatever it read off PiecewiseCallInputs is baked into the graph.

Parameters:
Return type:

PiecewiseOutput

set_slot(slot)[source]#
Parameters:

slot (int)

size_captures()[source]#

Warm every shape up and capture the largest once, thrown away; see AcceleratorGraphRunner.size_captures.

Return type:

dict[BucketKey, CaptureCost]

size_lent_captures(wanted)[source]#

Capture every shape a lender sized, thrown away, if wanted here or on another rank; see AcceleratorGraphRunner.size_lent_captures.

Parameters:

wanted (bool)

Return type:

dict[BucketKey, CaptureCost]

warmup_and_capture(planned)[source]#
Parameters:

planned (Collection[BucketKey] | None)

Return type:

None

class mstar.engine.accelerator_graph_runner.PiecewiseGraphData(graph, static_inputs, static_outputs, dummy_rids, shape, bucket)[source]#

Bases: object

One captured region, for one (bs, total_tokens) bucket.

Parameters:
bucket: BucketKey#
dummy_rids: list[int]#
graph: CapturedGraph#
shape: PiecewiseCaptureShape#
static_inputs: dict[str, Tensor]#
static_outputs: dict[str, Tensor]#
class mstar.engine.accelerator_graph_runner.PiecewiseGraphKey(bs, seq_len, slot)[source]#

Bases: NamedTuple

Parameters:
bs: int#

Alias for field number 0

seq_len: int#

Alias for field number 1

slot: int#

Alias for field number 2

class mstar.engine.accelerator_graph_runner.PiecewiseOutput(outputs, real_len)[source]#

Bases: object

Dict-like view over a captured region’s output buffers.

The runner replays into persistent buffers sized for the padded bucket; only the leading real_len rows are meaningful. Indexing and get return an owned CLONE of that leading slice — safe to keep past the next replay. get_view returns the same slice WITHOUT copying.

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

key (str)

get_view(key, default=None)[source]#

The leading real_len slice WITHOUT copying.

The result aliases the runner-owned static output buffer and is OVERWRITTEN by the next run of this runner. Nothing else writes there: the buffers live outside the graph memory pool, so replays of other regions sharing the node’s pool leave them alone. Read it within the same step; use get when you need something that outlives it.

Parameters:

key (str)

keys()[source]#
class mstar.engine.accelerator_graph_runner.WarmedRegion(run, static_inputs, dummy_rids, warm_outputs, peak)[source]#

Bases: NamedTuple

A region’s shape planned and run eagerly: what its capture reads.

Parameters:
dummy_rids: list[int]#

Alias for field number 2

peak: int#

Alias for field number 4

run: Callable[[], Any]#

Alias for field number 0

static_inputs: dict[str, Tensor]#

Alias for field number 1

warm_outputs: dict[str, Tensor]#

Alias for field number 3

class mstar.engine.accelerator_graph_runner.WarmedSpec(run, static_inputs, static_input_keys, dummy_rids, dummy_metadata, peak, returned)[source]#

Bases: NamedTuple

A spec prepared and run eagerly: what its capture reads.

Parameters:
dummy_metadata: dict[int, CurrentForwardPassInfo]#

Alias for field number 4

dummy_rids: list[int]#

Alias for field number 3

peak: int#

Alias for field number 5

returned: int#

Alias for field number 6

run: Callable[[], Any]#

Alias for field number 0

static_input_keys: tuple[str, ...]#

Alias for field number 2

static_inputs: dict[str, Any]#

Alias for field number 1

mstar.engine.accelerator_graph_runner.agree_across_ranks(comm_group, values, device, reduce)[source]#

Reduce each of values over every rank of the joint group with reduce. torch.amin ANDs flags and torch.amax takes the largest.

Capture failure is per-rank: a rank that keeps a graph another rank dropped replays it while the other runs eager, which hangs as soon as the captured region holds a collective. Callers pass an order derived from the configs, which are identical on every rank.

Parameters:
Return type:

list[int]

mstar.engine.accelerator_graph_runner.autocast_scope(dtype, device_type)[source]#

A forward’s autocast scope; None (disable_autocast) runs the submodule in its own dtype and shuts out any ambient autocast.

Parameters:

  • dtype (dtype | None)

  • device_type (str)

mstar.engine.accelerator_graph_runner.blocked_capture_resources(step, resources)[source]#

Return resource keys that prevent capture of the declared step.

Parameters:

resources (Mapping[str, Resource])

Return type:

list[str]

mstar.engine.accelerator_graph_runner.capture_into_graph(run, pool, device, autocast_dtype, keep_graph=False)[source]#

Capture run into a new CUDA graph allocated from pool.

Returns the graph and what run returned. A capture that fails part way leaves two things behind that would sink every later capture on the same pool: the caching allocator keeps routing this thread’s allocations into the pool (the next capture_begin fails with “already recording to mempool_id”), and torch.cuda.graph never switches the thread back off its capture stream. Undo both before re-raising, so the buckets after a failed one still capture and the thread keeps running on the stream it came in on.

A torch.compile’d forward that would recompile inside the capture is failed up front with the guard that broke: dynamo saves the CUDA RNG state before it compiles, which CUDA refuses during capture, so such a capture cannot succeed and would otherwise die with an unrelated error.

Parameters:
mstar.engine.accelerator_graph_runner.capture_with_static_outputs(run, warm_outputs, pool, device, autocast_dtype, keep_graph=False)[source]#

Capture run so its outputs land in buffers allocated OUTSIDE pool.

warm_outputs is what run returned on a warm-up call ({name: Tensor}); it fixes the shapes and dtypes. The captured graph ends with a device-to-device copy of each output into a buffer allocated here, before the capture, from the ordinary allocator.

Why not hand out the graph’s own output tensors: blocks a graph frees while it is being captured go back to the pool and are handed to the graphs captured after it, so an output that lived in the pool could sit exactly where an earlier-captured graph keeps its scratch and be overwritten by that graph’s next replay. That is harmless while a pool belongs to one region whose graphs replay one at a time and whose outputs are consumed before the next replay, but not once a node’s regions share one pool and replay in data-dependent order. Buffers outside the pool stay valid until this graph’s own next replay, whatever else the node runs in between.

Parameters:
mstar.engine.accelerator_graph_runner.compile_kwargs(config_mode)[source]#

torch.compile’s arguments for a capture config whose compile_mode is config_mode, which the config checked when it was built.

Parameters:

config_mode (str | None)

Return type:

dict[str, Any]

mstar.engine.accelerator_graph_runner.dummy_metadata(rids, graph_walk)[source]#

Stand-in request info for padding rows, which have no real request.

Parameters:
Return type:

dict[int, CurrentForwardPassInfo]

mstar.engine.accelerator_graph_runner.dummy_rid_handle(name)[source]#
Parameters:

name (str)

Return type:

int

mstar.engine.accelerator_graph_runner.dummy_rid_name(handle)[source]#

The name a padding-row handle was minted from. Debug/logging only.

Parameters:

handle (int)

Return type:

str

mstar.engine.accelerator_graph_runner.fail_if_graphs_required(missing)[source]#

MSTAR_REQUIRE_CUDA_GRAPHS=1 turns a dropped bucket into a startup failure, for deployments that would rather not come up than serve the eager path at 10-20x the latency.

Parameters:

missing (list[str])

Return type:

None

mstar.engine.accelerator_graph_runner.plan_captures(pools, free, reusable=None)[source]#

The cheapest buckets across every pool on a device, added while the largest eager step still fits in free beside them.

A pool’s graphs, captured largest first, reuse its memory: it holds what every graph keeps, plus the most any one graph needs beyond that. So a bucket adds its slots’ kept bytes, and whatever it needs past the pool’s most so far. That is a model; the boot log compares it with what capture took. A bucket also adds its slots’ executable memory, less the memory its graphs reuse from the walk’s copies. A walk’s copies pay for its planned graphs’ executable memory once, across all its buckets. reusable bounds the reuse summed over all buckets, None for no bound.

Parameters:
Return type:

CapturePlan

Contents