mstar.engine.resources.kv.ring.cache

mstar.engine.resources.kv.ring.cache#

Classes

LayerRingCache(num_sessions, n_kv_heads, ...)

One attention layer's ring: ring_frames frame slots of history plus one scratch frame at the tail, times num_sessions resident sessions.
class mstar.engine.resources.kv.ring.cache.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]