mstar.engine.resources.kv.ring

mstar.engine.resources.kv.ring#

Ring KV storage: a fixed horizon of frame slots, overwritten in place.

The storage strategy Waypoint declares. KVSpec.resource_class dispatches here on seeing a RingKVConfig, so a paged model never imports it and this package never imports FlashInfer.

class mstar.engine.resources.kv.ring.LayerRingCache(num_sessions, n_kv_heads, ring_frames, ring_buckets, d_head, tokens_per_frame, pinned_dilation, dtype, device)[source]#

Bases: object

One attention layer’s ring: ring_frames frame slots of history plus one scratch frame at the tail, times num_sessions resident sessions.

Storage is a single [2, 1, H_kv, num_sessions * capacity, D] tensor so that a commit is one index_copy_ for K and V together and the read is one unbind(0) into two views – no copy on the read path.

capacity is ONE session’s token slots; total_slots is the token-dim length of the buffer.

Parameters:
property memory_bytes: int#

Resident bytes of KV storage, all sessions (the bool/index buffers are noise).

reset(session_idx)[source]#

Drop ONE session’s state and re-arm its scratch tail.

Parameters:

session_idx (int)

Return type:

None

session_span(session_idx)[source]#

[lo, hi) token slots owned by session_idx.

A host int here, unlike everywhere on the forward path: the two callers below are host-side lifecycle (a rollout ending, capture tearing down), never inside a captured region.

Parameters:

session_idx (int)

Return type:

tuple[int, int]

upsert(kv, frame_pos, commit, session_idx, *, build_visibility=True)[source]#

kv is [2, 1, H_kv, B*tokens_per_frame, D], one frame for each of B sessions; commit writes the frame into its ring slot; without it the frame lands in that session’s scratch tail only, visible to itself and to nothing later.

Returns (k, v, visible)

Parameters:
Return type:

tuple[Tensor, Tensor, Tensor]

class mstar.engine.resources.kv.ring.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]

Modules

cache

manager

The paged KVManager reserves pages per step and hands them back; this one reserves nothing.