mstar.engine.chunking

mstar.engine.chunking#

Chunked prefill’s state on the engine.

Per (node, walk): the model’s chunking policy and the token budget a step runs under, after the serving config’s overrides. Per row mid-prefill: its whole prepared input, which every chunk is cut from. And the count of steps that ran over their budget, for the warning.

Classes

ChunkRange(start, end, total, stride)

The tokens [start, end) a chunked row runs this step, of the total its whole input was measured at.

ChunkingManagement(submodules)

class mstar.engine.chunking.ChunkRange(start, end, total, stride)[source]#

Bases: NamedTuple

The tokens [start, end) a chunked row runs this step, of the total its whole input was measured at.

stride is how far the chunk moves the row’s chunk Loop, which counts the tokens run up to total (GraphRuntime.set_loop_caps): end for a prompt’s first chunk, whose Loop starts at 0 even when a cached prefix moves start on, and end - start for each later one.

Parameters:
end: int#

Alias for field number 1

property is_final: bool#
start: int#

Alias for field number 0

stride: int#

Alias for field number 3

total: int#

Alias for field number 2

class mstar.engine.chunking.ChunkingManagement(submodules)[source]#

Bases: object

Parameters:

submodules (Mapping[str, SubmoduleManagement])

check_budget(batch, node_inputs)[source]#

Warn when a step runs over the budget it was scheduled under: that of the combined walk it ran in, if any. A row too long to cut runs whole, and a speculation merge sizes its continuing rows by estimate, before their inputs exist. Logged on the 1st, 2nd, 4th, … overrun per walk.

Parameters:
Return type:

None

chunkable(node_name, graph_walk)[source]#

The submodule’s say, unless the serving config turned the walk’s chunking off (a null budget) or the worker chunks nothing on the node.

Parameters:

  • node_name (str)

  • graph_walk (str)

Return type:

bool

forget_request(rid)[source]#
Parameters:

rid (int)

Return type:

None

max_batch_tokens(node_name, graph_walk)[source]#

The walk’s budget, unless the serving config overrides it or the worker chunks nothing on the node.

Parameters:

  • node_name (str)

  • graph_walk (str)

Return type:

int | None

min_chunk_tokens(budget)[source]#

The fewest tokens the scheduler cuts a row to under budget, or the whole budget when that is smaller. Without a floor, decode rows that fill the budget leave a prompt a token or two a step.

Parameters:

budget (int)

Return type:

int

policy(node_name, graph_walk)[source]#

The model’s, or NO_CHUNKING where it declares none.

Declared on the model rather than the submodule because the graph build needs it too: a chunkable walk runs the node inside a Loop, and get_graph_walk_graphs is read long before any submodule loads.

Parameters:

  • node_name (str)

  • graph_walk (str)

Return type:

ChunkingPolicy

prepared: dict[tuple[int, str], NodeInputs]#
release(rid, node_name)[source]#

Drop a row’s whole input once its last chunk lands.

Parameters:
Return type:

None

set_budgets(overrides, walks, combined_walk_of=None, override_mixed=(), no_chunking=(), min_chunk_tokens=DEFAULT_MIN_CHUNK_TOKENS)[source]#

Apply a serving config’s max_batch_tokens and min_chunk_tokens to this engine’s nodes.

Overrides the budgets set_policies declared, so it runs after it.

Per node, one budget for every walk the model budgets, or a mapping of walk -> budget; null turns chunking off. Only walks the model budgets can be overridden: elsewhere the setting would do nothing.

A walk inside a combined walk runs in the combined walk’s steps, under the combined walk’s budget, so its own budget applies there only on a node in override_mixed (can_override_mixed_max_tokens), where it caps that walk’s rows within the step. Elsewhere a per-walk override of such a walk would do nothing, so it is refused; the combined walk itself can be named instead.

no_chunking names nodes this worker chunks nowhere (SP, which chunked prefill does not support yet). A budget set on one is still validated but has no effect, since the deployment, not the config, is what turned chunking off there.

Parameters:
Return type:

None

set_policies(policies)[source]#

The model’s chunking_policy_of(), bound at load.

Separate from set_budgets: the policy is the model’s and fixed, the budget is the deployment’s and layered on top.

Parameters:

policies (Mapping[tuple[str, str], ChunkingPolicy])

Return type:

None

walk_caps(node_name, graph_walk)[source]#

Under a combined walk on a node with can_override_mixed_max_tokens: walk -> the most tokens that walk’s rows may run in one step, for the walks with a budget of their own. None otherwise.

Parameters:

  • node_name (str)

  • graph_walk (str)

Return type:

dict[str, int] | None