mstar.engine.resources.kv.bounded.manager

mstar.engine.resources.kv.bounded.manager#

A pool of fixed-size per-request KV slots with sink + window retention.

Unlike the paged KVManager there is nothing to reserve per step and nothing to release: a request takes a slot when it first runs and keeps it, the window is a ring overwritten in place, and a step’s layout is a small integer table computed on the host from positions the step declares.

Classes

BoundedKVManager(config, device[, kv_dtype])

BoundedPlan(table, rows, max_retained, ...)

class mstar.engine.resources.kv.bounded.manager.BoundedKVManager(config, device, kv_dtype=None)[source]#

Bases: Resource

Parameters:
admit(step, ctx)[source]#

Reserve space for the given step. In the case where requests in a batch must be executed sequentially, this may be called for all requests in a loop before the per-request plan -> forward -> commit cycle.

Parameters:
Return type:

AdmitOutcome

attend(layer_idx, q, k, v, source=None, plan=None, rel_bias=None, ieee=False, row_source=None)[source]#

Attention for q, k, v [B * rows, H, T, D] over each row’s retained keys (source [sources, rows, H, L, 2D] is this layer’s source prefixes, each row’s picked by its StreamPosition.source) and this step’s own, then this step’s keys into the slots. [B * rows, T, H, D]. plan defaults to the step’s. rel_bias is a per-distance score term (kernels.bounded_attention), for relative-position attention; with a per-query window (BoundedKVConfig.query_window) it is indexed by distance, [B * rows, H, T, window + 1]. Rows whose source is -1 read theirs from row_source = (storage [slots, rows, H, L, 2D], slots [>= B]), e.g. another resource’s per-request entries.

Parameters:
Return type:

Tensor

classmethod build(spec, info)[source]#
Parameters:
Return type:

BoundedKVManager

build_accelerator_graph_buffers(slots, max_bs, max_seq_len)[source]#

Size whatever the captured replays will read.

Called once per runner that captures against this node — the whole forward’s, and one per piecewise region — so it must tolerate repeated calls: grow to the largest shape asked for, never clobber what an earlier call already sized.

Parameters:
Return type:

None

cleanup()[source]#
commit(step, ctx)[source]#

record step consumption

Parameters:
Return type:

None

property current: BoundedPlan#
ingest_request(rid, overrides=None)[source]#

Must be idempotent: the conductor sends one NewRequest per partition, so a worker serving several partitions of a request ingests it once for each, possibly after a node has already run a step for it.

Parameters:
layer(layer_idx)[source]#

One layer’s slots, [max_slots, rows, H, cap, 2D].

Parameters:

layer_idx (int)

Return type:

Tensor

plan(step, ctx)[source]#

ret is immutable and opaque to runner; only gives to ctx.plan_results

Parameters:
Return type:

BoundedPlan

post_warmup_validate()[source]#

For, e.g., the KV cache to check that num_free_pages is identical across TP ranks after cuda graph capture.

Raises an error (fails loudly) if invalid.

remove_request(rid)[source]#
Parameters:

rid (str)

reset_request(rid, free=False)[source]#

Nothing to zero: a slot’s contents are only read where a later position says they were written.

Parameters:
rollback_admit(step, ctx)[source]#

Free the slots this step’s admit took fresh; a slot a request already held stays with it.

Parameters:
Return type:

None

class mstar.engine.resources.kv.bounded.manager.BoundedPlan(table, rows, max_retained, max_span, reads_source, reads_row_source)[source]#

Bases: NamedTuple

Parameters:
max_retained: int#

Alias for field number 2

max_span: int#

Alias for field number 3

reads_row_source: bool#

Alias for field number 5

reads_source: bool#

Alias for field number 4

retained()[source]#

Each row’s retained keys before its step, [rows] int32 on the device.

Return type:

Tensor

rows: tuple[tuple[int, RowLayout, int], ...]#

Alias for field number 1

table: Tensor#

Alias for field number 0