mstar.engine.cuda_graph_runner#
Functions
|
AND each flag across every rank of the joint group. |
|
A forward's autocast scope; |
|
Stand-in request info for padding rows, which have no real request. |
Classes
|
The captured slots for one (walk, cg_key_info, bs, num_tokens). |
|
|
|
One captured graph + the buffers its replay reads and writes. |
|
The padding rows captured replays pad onto. |
|
Captures one inner callable of a submodule's forward as a CUDA graph. |
|
One captured region, for one (bs, total_tokens) bucket. |
|
Dict-like view over a captured region's output buffers. |
- class mstar.engine.cuda_graph_runner.CudaGraphBucket(config, config_idx, slots=<factory>)[source]#
Bases:
objectThe captured slots for one (walk, cg_key_info, bs, num_tokens).
- Parameters:
config (CudaGraphConfig)
config_idx (int)
slots (list[CudaGraphSlot])
- config: CudaGraphConfig#
- slots: list[CudaGraphSlot]#
- class mstar.engine.cuda_graph_runner.CudaGraphRunner(submodule_name, submodule, resources, step_runner, device, autocast_dtype, joint_comm_group, enable_nvtx=False)[source]#
Bases:
object- Parameters:
submodule_name (str)
submodule (NodeSubmodule)
step_runner (StepRunner)
device (device)
autocast_dtype (dtype | None)
joint_comm_group (JointGroups)
enable_nvtx (bool)
- CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32, 64]#
- NUM_SLOTS = 2#
- NUM_WARMUP = 2#
- property any_graphs#
- 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_stepand nothing else — the list is shared, so a caller that mutates it or passes it topreprocesscorrupts every later step on this bucket.- Parameters:
lease (SlotLease)
- Return type:
- lease_slot(graph_walk, bs, num_tokens=None, cg_key_info=None, slot=None)[source]#
Pick the bucket and double-buffer slot an upcoming step replays on.
num_tokens=Nonemeans the batch’s inputs aren’t built yet, so the bucket comes from the batched-capture search instead.slot=Noneadvances 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.
- 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.
- 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:
lease (SlotLease)
inputs (list[NodeInputs])
- Return type:
- 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:
- release(lease, real_bs)[source]#
Return the padding rows to their at-rest state after a step.
Their pages stay resident (
free=False), so the next step’s plan for this slot allocates nothing for the tail.
- 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_eventis 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_eventreleases 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++.
- 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 (
bsrows 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 forprepare_inputsand go throughselect_bucket.This is what lets pre-plan run before inputs exist. Lifting the restriction means making
prepare_inputssafe to run ahead of the forward — seeEngine.pre_plan_for_batch.
- 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.
- 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.
- class mstar.engine.cuda_graph_runner.CudaGraphSlot(graph, static_inputs, static_input_keys, static_outputs, dummy_rids, dummy_metadata, config_idx)[source]#
Bases:
objectOne 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:
- dummy_metadata: dict[str, CurrentForwardPassInfo]#
- class mstar.engine.cuda_graph_runner.DummyRowPool(prefix, step_runner, resources)[source]#
Bases:
objectThe 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:
prefix (str)
step_runner (StepRunner)
- class mstar.engine.cuda_graph_runner.PiecewiseCudaGraphRunner(label, config, resources, step_runner, device, autocast_dtype, joint_comm_group=None, node_name=None)[source]#
Bases:
objectCaptures one inner callable of a submodule’s forward as a CUDA graph.
Where
CudaGraphRunnerreplays a wholeforward_batchedunder engine control, this captures a SUB-REGION — a transformer block loop, say — while the surrounding preamble stays eager and the submodule invokesrunitself. 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:
label (str)
config (PiecewiseCudaGraphConfig)
step_runner (StepRunner)
device (device)
autocast_dtype (dtype | None)
joint_comm_group (JointGroups | None)
node_name (str | None)
- CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32, 64]#
- NUM_WARMUP = 2#
- SLOT = 0#
- lease_slot(batch_size, total_tokens=None)[source]#
The bucket this batch would replay on, asked before the forward.
Same question
runsettles per call, answered early so the outerdeclare_stepcan plan the region’s resources against it.
- prepare_for_capture()[source]#
Claim this region’s static buffers; see
CudaGraphRunner.- Return type:
- 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
PiecewiseCallInputsis baked into the graph.
- class mstar.engine.cuda_graph_runner.PiecewiseGraphData(graph, static_inputs, static_outputs, dummy_rids, shape, bucket)[source]#
Bases:
objectOne captured region, for one (bs, total_tokens) bucket.
- Parameters:
- shape: PiecewiseCaptureShape#
- class mstar.engine.cuda_graph_runner.PiecewiseOutput(outputs, real_len)[source]#
Bases:
objectDict-like view over a captured region’s output buffers.
The runner replays into persistent buffers sized for the padded bucket; only the leading
real_lenrows are meaningful. Indexing andgetreturn an owned CLONE of that leading slice — safe to keep past the next replay.get_viewreturns the same slice WITHOUT copying.
- 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.