mstar.engine.accelerator_graph_config

Functions

check_compile_mode(mode)	A capture config's compile_mode, refused at build rather than at capture, where the per-capture fallback would serve the bucket eager.
decode_rows_then_prompt(total_tokens, bs)	bs - 1 one-token rows and a prompt row of the rest: a mixed step's layout, for kernels that split rows by span.
distribute_tokens(total_tokens, bs)	Split total_tokens across bs requests, remainder on the first.

Classes

AcceleratorGraphConfig(capture_graph_walk[, ...])
AcceleratorGraphConfigType(*values)
BatchedAcceleratorGraphConfig(...[, ...])
PackedAcceleratorGraphConfig(...[, ...])
PiecewiseAcceleratorGraphConfig(*, ...[, ...])	One inner callable of a submodule's forward, captured on its own.
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)
PiecewisePackedConfig(*, capture_fn, ...[, ...])	Packed variable-length capture: static input [total_tokens, D].

class mstar.engine.accelerator_graph_config.AcceleratorGraphConfig(capture_graph_walk, replay_graph_walks=None, additional_key_info=None, compile=True, compile_mode=None, capture_batch_sizes=None, capture_forward_method='forward_batched', caps_eager_batch_size=True, input_seq_dims=None)

Bases: ABC

Parameters:

• capture_graph_walk (str)

• replay_graph_walks (list[str] | None)

• additional_key_info (Any | None)

• compile (bool)

• compile_mode (str | None)

• capture_batch_sizes (list[int] | None)

• capture_forward_method (str)

• caps_eager_batch_size (bool)

• input_seq_dims (dict[str, int] | None)

abstractmethod get_config_type()

Return type:

AcceleratorGraphConfigType

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.accelerator_graph_config.AcceleratorGraphConfigType(*values)

Bases: Enum

BASIC_BATCHED = 'basic_batched'

FLASH_INFER_PACKED = 'flash_infer_packed'

class mstar.engine.accelerator_graph_config.BatchedAcceleratorGraphConfig(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, compile_mode=None, input_seq_dims=None)

Bases: AcceleratorGraphConfig

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)

• compile_mode (str | None)

• input_seq_dims (dict[str, int] | None)

get_config_type()

Return type:

AcceleratorGraphConfigType

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.accelerator_graph_config.PackedAcceleratorGraphConfig(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, total_tokens_multiplier=1, compile_mode=None, input_seq_dims=None, token_layout=None)

Bases: AcceleratorGraphConfig

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)

• total_tokens_multiplier (int)

• compile_mode (str | None)

• input_seq_dims (dict[str, int] | None)

• token_layout (Callable[[int, int], list[int]] | None)

get_config_type()

Return type:

AcceleratorGraphConfigType

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.accelerator_graph_config.PiecewiseAcceleratorGraphConfig(*, capture_fn, make_static_inputs, declare_step=None, lease_before_step=False, can_lease=None, forward_kwargs=<factory>, capture_batch_sizes=None, compile=False, compile_mode=None)

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)

• can_lease (Callable[[list[NodeInputs]], bool] | None)

• forward_kwargs (dict[str, Any])

• capture_batch_sizes (list[int] | None)

• compile (bool)

• compile_mode (str | None)

can_lease: Callable[[list[NodeInputs]], bool] | None = None

capture_batch_sizes: list[int] | None = None

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

compile: bool = False

compile_mode: str | None = None

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.accelerator_graph_config.PiecewiseBatchedConfig(*, capture_fn, make_static_inputs, declare_step=None, lease_before_step=False, can_lease=None, forward_kwargs=<factory>, capture_batch_sizes=None, compile=False, compile_mode=None, seq_len)

Bases: PiecewiseAcceleratorGraphConfig

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)

• can_lease (Callable[[list[NodeInputs]], bool] | None)

• forward_kwargs (dict[str, Any])

• capture_batch_sizes (list[int] | None)

• compile (bool)

• compile_mode (str | None)

• 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.accelerator_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.accelerator_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.accelerator_graph_config.PiecewiseConfigType(*values)

Bases: Enum

BATCHED = 'batched'

PACKED = 'packed'

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

Bases: PiecewiseAcceleratorGraphConfig

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)

• can_lease (Callable[[list[NodeInputs]], bool] | None)

• forward_kwargs (dict[str, Any])

• capture_batch_sizes (list[int] | None)

• compile (bool)

• compile_mode (str | None)

• 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.accelerator_graph_config.check_compile_mode(mode)

A capture config’s compile_mode, refused at build rather than at capture, where the per-capture fallback would serve the bucket eager.

Parameters:

mode (str | None)

Return type:

str | None

mstar.engine.accelerator_graph_config.decode_rows_then_prompt(total_tokens, bs)

bs - 1 one-token rows and a prompt row of the rest: a mixed step’s layout, for kernels that split rows by span. distribute_tokens when there are fewer tokens than rows.

Parameters:

• total_tokens (int)

• bs (int)

Return type:

list[int]

mstar.engine.accelerator_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]