mstar.engine.cuda_graph_config

Functions

distribute_tokens(total_tokens, bs)

Split total_tokens across bs requests, remainder on the first.

Classes

BatchedCudaGraphConfig(capture_graph_walk, ...)
CudaGraphConfig(capture_graph_walk[, ...])
CudaGraphConfigType(*values)
PackedCudaGraphConfig(capture_graph_walk, ...)
PiecewiseBatchedConfig(*, capture_fn, ...[, ...])	Equal-length batched capture: static input [bs, seq_len, D].
PiecewiseCallInputs(static_inputs, engine_inputs)	Everything a captured region is handed, at capture and at replay alike.
PiecewiseCaptureShape(bs, seq_lens, total_tokens)	One piecewise capture bucket.
PiecewiseConfigType(*values)
PiecewiseCudaGraphConfig(*, capture_fn, ...)	One inner callable of a submodule's forward, captured on its own.
PiecewisePackedConfig(*, capture_fn, ...[, ...])	Packed variable-length capture: static input [total_tokens, D].

class mstar.engine.cuda_graph_config.BatchedCudaGraphConfig(capture_graph_walk, single_request_inputs, replay_graph_walks=None, additional_key_info=None, compile=True, capture_batch_sizes=None, capture_forward_method='forward_batched', caps_eager_batch_size=True, total_tokens_multiplier=1)

Bases: CudaGraphConfig

Parameters:

• capture_graph_walk (str)

• single_request_inputs (NodeInputs)

• replay_graph_walks (list[str] | None)

• additional_key_info (Any | None)

• compile (bool)

• capture_batch_sizes (list[int] | None)

• capture_forward_method (str)

• caps_eager_batch_size (bool)

• total_tokens_multiplier (int)

get_config_type()

Return type:

CudaGraphConfigType

get_node_inputs(bs, num_tokens)

Parameters:

• bs (int)

• num_tokens (int)

get_total_tokens(bs)

Parameters:

bs (int)

Return type:

list[int]

class mstar.engine.cuda_graph_config.CudaGraphConfig(capture_graph_walk, replay_graph_walks=None, additional_key_info=None, compile=True, capture_batch_sizes=None, capture_forward_method='forward_batched', caps_eager_batch_size=True)

Bases: ABC

Parameters:

• capture_graph_walk (str)

• replay_graph_walks (list[str] | None)

• additional_key_info (Any | None)

• compile (bool)

• capture_batch_sizes (list[int] | None)

• capture_forward_method (str)

• caps_eager_batch_size (bool)

abstractmethod get_config_type()

Return type:

CudaGraphConfigType

abstractmethod get_node_inputs(bs, num_tokens)

Parameters:

• bs (int)

• num_tokens (int)

Return type:

list[NodeInputs]

abstractmethod get_total_tokens(bs)

Parameters:

bs (int)

Return type:

list[int]

class mstar.engine.cuda_graph_config.CudaGraphConfigType(*values)

Bases: Enum

BASIC_BATCHED = 'basic_batched'

FLASH_INFER_PACKED = 'flash_infer_packed'

class mstar.engine.cuda_graph_config.PackedCudaGraphConfig(capture_graph_walk, capture_token_lengths, make_node_input, replay_graph_walks=None, additional_key_info=None, compile=True, capture_batch_sizes=None, capture_forward_method='forward_batched', caps_eager_batch_size=True)

Bases: CudaGraphConfig

Parameters:

• capture_graph_walk (str)

• capture_token_lengths (list[int])

• make_node_input (Callable[[int], NodeInputs])

• replay_graph_walks (list[str] | None)

• additional_key_info (Any | None)

• compile (bool)

• capture_batch_sizes (list[int] | None)

• capture_forward_method (str)

• caps_eager_batch_size (bool)

get_config_type()

Return type:

CudaGraphConfigType

get_node_inputs(bs, num_tokens)

Parameters:

• bs (int)

• num_tokens (int)

get_total_tokens(bs)

Parameters:

bs (int)

Return type:

list[int]

class mstar.engine.cuda_graph_config.PiecewiseBatchedConfig(*, capture_fn, make_static_inputs, declare_step=None, lease_before_step=False, forward_kwargs=<factory>, capture_batch_sizes=None, compile=False, seq_len)

Bases: PiecewiseCudaGraphConfig

Equal-length batched capture: static input [bs, seq_len, D].

Parameters:

• capture_fn (Callable[[PiecewiseCallInputs], dict[str, Tensor]])

• make_static_inputs (Callable[[PiecewiseCaptureShape], dict[str, Tensor]])

• declare_step (Callable[[list[str], list[int]], SubmoduleStep | None] | None)

• lease_before_step (bool)

• forward_kwargs (dict[str, Any])

• capture_batch_sizes (list[int] | None)

• compile (bool)

• seq_len (int)

get_capture_shapes(batch_sizes)

The (bs, seq_lens, total_tokens) buckets to capture.

batch_sizes is the resolved list the runner iterates (capture_batch_sizes or the runner default).

Parameters:

batch_sizes (list[int])

Return type:

list[PiecewiseCaptureShape]

get_config_type()

Return type:

PiecewiseConfigType

replay_seq_lens(shape, seq_lens, real_bs)

Per-request lengths to declare at replay, padded to shape.bs.

Parameters:

• shape (PiecewiseCaptureShape)

• seq_lens (list[int] | None)

• real_bs (int)

Return type:

list[int]

seq_len: int

class mstar.engine.cuda_graph_config.PiecewiseCallInputs(static_inputs, engine_inputs, kwargs=<factory>)

Bases: object

Everything a captured region is handed, at capture and at replay alike.

Both calls go through here so a region can’t accidentally read something that only exists on one of the two paths.

Parameters:

• static_inputs (dict[str, Tensor])

• engine_inputs (ModelInputsFromEngine)

• kwargs (dict[str, Any])

engine_inputs: ModelInputsFromEngine

kwargs: dict[str, Any]

property resources: dict[str, Any]

static_inputs: dict[str, Tensor]

class mstar.engine.cuda_graph_config.PiecewiseCaptureShape(bs, seq_lens, total_tokens)

Bases: object

One piecewise capture bucket.

Handed to the static-input factory and the step declaration so both generalize across config types.

Parameters:

• bs (int)

• seq_lens (list[int])

• total_tokens (int)

bs: int

seq_lens: list[int]

total_tokens: int

class mstar.engine.cuda_graph_config.PiecewiseConfigType(*values)

Bases: Enum

BATCHED = 'batched'

PACKED = 'packed'

class mstar.engine.cuda_graph_config.PiecewiseCudaGraphConfig(*, capture_fn, make_static_inputs, declare_step=None, lease_before_step=False, forward_kwargs=<factory>, capture_batch_sizes=None, compile=False)

Bases: ABC

One inner callable of a submodule’s forward, captured on its own.

Unlike CudaGraphConfig, which describes a whole forward_batched the engine drives, this describes a SUB-REGION the submodule invokes itself while the surrounding preamble stays eager.

kw_only so subclasses can add required fields (e.g. seq_len) without colliding with the defaulted ones here.

Parameters:

• capture_fn (Callable[[PiecewiseCallInputs], dict[str, Tensor]])

• make_static_inputs (Callable[[PiecewiseCaptureShape], dict[str, Tensor]])

• declare_step (Callable[[list[str], list[int]], SubmoduleStep | None] | None)

• lease_before_step (bool)

• forward_kwargs (dict[str, Any])

• capture_batch_sizes (list[int] | None)

• compile (bool)

capture_batch_sizes: list[int] | None = None

capture_fn: Callable[[PiecewiseCallInputs], dict[str, Tensor]]

compile: bool = False

declare_step: Callable[[list[str], list[int]], SubmoduleStep | None] | None = None

forward_kwargs: dict[str, Any]

abstractmethod get_capture_shapes(batch_sizes)

The (bs, seq_lens, total_tokens) buckets to capture.

batch_sizes is the resolved list the runner iterates (capture_batch_sizes or the runner default).

Parameters:

batch_sizes (list[int])

Return type:

list[PiecewiseCaptureShape]

abstractmethod get_config_type()

Return type:

PiecewiseConfigType

lease_before_step: bool = False

make_static_inputs: Callable[[PiecewiseCaptureShape], dict[str, Tensor]]

abstractmethod replay_seq_lens(shape, seq_lens, real_bs)

Per-request lengths to declare at replay, padded to shape.bs.

Parameters:

• shape (PiecewiseCaptureShape)

• seq_lens (list[int] | None)

• real_bs (int)

Return type:

list[int]

class mstar.engine.cuda_graph_config.PiecewisePackedConfig(*, capture_fn, make_static_inputs, declare_step=None, lease_before_step=False, forward_kwargs=<factory>, capture_batch_sizes=None, compile=False, total_tokens)

Bases: PiecewiseCudaGraphConfig

Packed variable-length capture: static input [total_tokens, D].

One graph per (bs, token bucket). Each bucket is partitioned across bs requests for the capture-time plan; real per-request lengths arrive at replay.

Parameters:

• capture_fn (Callable[[PiecewiseCallInputs], dict[str, Tensor]])

• make_static_inputs (Callable[[PiecewiseCaptureShape], dict[str, Tensor]])

• declare_step (Callable[[list[str], list[int]], SubmoduleStep | None] | None)

• lease_before_step (bool)

• forward_kwargs (dict[str, Any])

• capture_batch_sizes (list[int] | None)

• compile (bool)

• total_tokens (list[int])

get_capture_shapes(batch_sizes)

The (bs, seq_lens, total_tokens) buckets to capture.

batch_sizes is the resolved list the runner iterates (capture_batch_sizes or the runner default).

Parameters:

batch_sizes (list[int])

Return type:

list[PiecewiseCaptureShape]

get_config_type()

Return type:

PiecewiseConfigType

replay_seq_lens(shape, seq_lens, real_bs)

Per-request lengths to declare at replay, padded to shape.bs.

Parameters:

• shape (PiecewiseCaptureShape)

• seq_lens (list[int] | None)

• real_bs (int)

Return type:

list[int]

total_tokens: list[int]

mstar.engine.cuda_graph_config.distribute_tokens(total_tokens, bs)

Split total_tokens across bs requests, remainder on the first.

Parameters:

• total_tokens (int)

• bs (int)

Return type:

list[int]