mstar.engine.resources.kv.ring.manager

mstar.engine.resources.kv.ring.manager#

The paged KVManager reserves pages per step and hands them back; this one reserves nothing. Every layer’s ring is allocated once at build and reused for the life of the process,.

Classes

RingKVManager(config, name, device[, dtype])

Per-layer ring caches holding num_sessions sessions, one per request.

RingPlan(request_ids, session_idx, frame_pos)

Host facts downstream resources need to stage this fixed ring view, one entry per row of the step batch, in ctx.request_ids order.
class mstar.engine.resources.kv.ring.manager.RingKVManager(config, name, device, dtype=torch.bfloat16)[source]#

Bases: AttentionResource

Per-layer ring caches holding num_sessions sessions, one per request.

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

Claim a session for this step’s request. A full pool is retried; a malformed step is refused terminally. Does not support eviction for now

Parameters:
Return type:

AdmitOutcome

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

RingKVManager

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

No-op: every buffer a replay touches was allocated at build, _static_session_idx included.

Parameters:
Return type:

None

capacity(layer_idx)[source]#

Token slots ONE session owns in layer_idx’s ring, scratch frame included. total_slots is the whole buffer.

Parameters:

layer_idx (int)

Return type:

int

cleanup()[source]#
Return type:

None

commit(step, ctx)[source]#

Record the frame each request just committed. Metadata only.

Parameters:
Return type:

None

get_state(rid)[source]#

Snapshot one request’s session. Cloned, so the caller can hold it across further rollout steps that mutate the rings in place.

Parameters:

rid (str)

Return type:

dict

ingest_request(rid, overrides=None)[source]#

Register rid. Deliberately does not claim a session.

Parameters:
Return type:

None

load_state(rid, state)[source]#

Restore one session’s contents into the existing allocation.

Parameters:
Return type:

None

memory_bytes()[source]#

Total resident ring bytes across all layers and all sessions.

Return type:

int

property num_sessions: int#

How many requests can hold a session here at once.

plan(step, ctx)[source]#

Stage the padded batch’s session indices, one per row. Ring addresses stay in the graph.

A replay pads the batch to its capture bucket with dummy rids that hold no session. Those rows still write and attend (their output is dropped downstream), so every one is parked on _padding_session – the shared scratch session outside the resident pool. The flex mask scopes each row to its own session, so a padding row’s read and its commit both stay on that session and never reach a real request’s history; no request is ever admitted to it, so the tail it leaves behind is never read.

Parameters:
Return type:

RingPlan

post_warmup_validate()[source]#

Capture must leave every session exactly as it found it.

Return type:

None

remove_request(rid)[source]#

The request is gone: drop its session, its claim, and its registration.

Parameters:

rid (str)

Return type:

None

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

Zero rid’s session and release its claim. free is ignored, there is no physical allocation to hand back.

Parameters:
Return type:

None

session_of(rid)[source]#

rid’s session index, or None if it holds none. Host-side introspection; the forward reads the staged tensor, never this.

Parameters:

rid (str)

Return type:

int | None

property tokens_per_frame: int#
total_slots(layer_idx)[source]#

Token slots in layer_idx’s buffer across every session – the length of the KV view upsert returns and of its visibility row.

Parameters:

layer_idx (int)

Return type:

int

upsert(k, v, layer_idx, frame_pos, *, commit, build_visibility=True)[source]#

Write one step’s K/V for layer_idx and return what to attend to.

k/v are [B, H_kv, tokens_per_frame, D], one frame per resident session in the step batch. k is already RoPE’d and RMS-normed and v is post value-residual lerp: the cache stores post-RoPE keys, so replayed history is never re-rotated. Returns (k_all, v_all, visible), the first two spanning the whole buffer, every resident session and the third a [total_slots] bool row that is False everywhere outside the calling request’s own session.

frame_pos and commit are both arguments rather than resource state, for the same reason. frame_pos is the [B] int64 ring clock – not a slot id – and alone determines the slot written and the visibility row;

Parameters:
Return type:

tuple[Tensor, Tensor, Tensor]

class mstar.engine.resources.kv.ring.manager.RingPlan(request_ids, session_idx, frame_pos)[source]#

Bases: NamedTuple

Host facts downstream resources need to stage this fixed ring view, one entry per row of the step batch, in ctx.request_ids order.

Parameters:
frame_pos: tuple[int, ...]#

Alias for field number 2

request_ids: tuple[str, ...]#

Alias for field number 0

session_idx: tuple[int, ...]#

Alias for field number 1