mstar.engine.cuda_graph_runner

Contents

mstar.engine.cuda_graph_runner#

Functions

agree_across_ranks(comm_group, flags, device)

AND each flag across every rank of the joint group.

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.

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.

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.

Classes

CudaGraphBucket(config, config_idx[, slots])

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

CudaGraphRunner(submodule_name, submodule, ...)

CudaGraphSlot(graph, static_inputs, ...)

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

DummyRowPool(prefix, step_runner, resources)

The padding rows captured replays pad onto.

PiecewiseCudaGraphRunner(label, config, ...)

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.
class mstar.engine.cuda_graph_runner.CudaGraphBucket(config, config_idx, slots=<factory>)[source]#

Bases: object

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

Parameters:
config: CudaGraphConfig#
config_idx: int#
slots: list[CudaGraphSlot]#
class mstar.engine.cuda_graph_runner.CudaGraphRunner(submodule_name, submodule, resources, step_runner, device, autocast_dtype, joint_comm_group, num_slots, enable_nvtx=False)[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:

CudaGraphConfig

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

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

slot_for(lease)[source]#
Parameters:

lease (SlotLease)

Return type:

CudaGraphSlot

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

Capture graphs for all configs and batch sizes.

class mstar.engine.cuda_graph_runner.CudaGraphSlot(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: CUDAGraph#
static_input_keys: tuple[str, ...]#
static_inputs: dict[str, Any]#
static_outputs: BatchedModelOutput#
class mstar.engine.cuda_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.cuda_graph_runner.PiecewiseCudaGraphRunner(label, config, resources, step_runner, device, autocast_dtype, num_slots, joint_comm_group=None, node_name=None, memory_pool=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_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)

warmup_and_capture()[source]#
Return type:

None

class mstar.engine.cuda_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: CUDAGraph#
shape: PiecewiseCaptureShape#
static_inputs: dict[str, Tensor]#
static_outputs: dict[str, Tensor]#
class mstar.engine.cuda_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.cuda_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]#
mstar.engine.cuda_graph_runner.agree_across_ranks(comm_group, flags, device)[source]#

AND each flag across every rank of the joint group.

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

mstar.engine.cuda_graph_runner.autocast_scope(dtype, device_type='cuda')[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.cuda_graph_runner.capture_into_graph(run, pool, device, autocast_dtype)[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.

mstar.engine.cuda_graph_runner.capture_with_static_outputs(run, warm_outputs, pool, device, autocast_dtype)[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.

mstar.engine.cuda_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.cuda_graph_runner.dummy_rid_handle(name)[source]#
Parameters:

name (str)

Return type:

int

mstar.engine.cuda_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.cuda_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

Contents