mstar.engine.resources.kv.manager

mstar.engine.resources.kv.manager#

Classes

`AllocResult`([success, error])
`CacheStream`([page_indices, stored_len, ...])	(request, label) cache stream metadata
`ClaimedStream`(label, pages, generation, ...)	A stream an in-progress offload has taken ownership of, and the state its host copy was made from.
`KVManager`(cfg, name, joint_comm_group, ...)
`KVPlanState`(token_to_page, token_to_cache[, ...])
`KVSequenceInfo`(seq_len, ...)
`PageArena`(kv_cache, allocator)	physical storage and free list management
`PublishedKVInfo`(info, dict[str, ...)
`RetentionPolicy`(context_budget)	fifo retention of context_budget

class mstar.engine.resources.kv.manager.AllocResult(success: bool = True, error: mstar.engine.resources.step.AdmitFailedReason | None = None)[source]#

Bases: object

Parameters:

success (bool)
error (AdmitFailedReason | None)

error: AdmitFailedReason | None = None#

success: bool = True#

class mstar.engine.resources.kv.manager.CacheStream(page_indices=<factory>, stored_len=0, position=0, released=0, retention=None, read_pending=False, read_future=None, read_error=None, offloaded=False, generation=0, step_in_flight=False)[source]#

Bases: object

(request, label) cache stream metadata

Parameters:

page_indices (list[int])
stored_len (int)
position (int)
released (int)
retention (RetentionPolicy | None)
read_pending (bool)
read_future (Future | None)
read_error (BaseException | None)
offloaded (bool)
generation (int)
step_in_flight (bool)

generation: int = 0#

offloaded: bool = False#

page_indices: list[int]#

position: int = 0#

read_error: BaseException | None = None#

read_future: Future | None = None#

read_pending: bool = False#

released: int = 0#

reset(freed=False)[source]#

Parameters:: freed (bool)

retention: RetentionPolicy | None = None#

step_in_flight: bool = False#

stored_len: int = 0#

class mstar.engine.resources.kv.manager.ClaimedStream(label, pages, generation, stored_len, position, released)[source]#

Bases: object

A stream an in-progress offload has taken ownership of, and the state its host copy was made from.

Parameters:

label (str)
pages (list[int])
generation (int)
stored_len (int)
position (int)
released (int)

generation: int#

label: str#

pages: list[int]#

position: int#

released: int#

stored_len: int#

class mstar.engine.resources.kv.manager.KVManager(cfg, name, joint_comm_group, transfer_engine_info, device, dtype=torch.bfloat16)[source]#

Bases: AttentionResource

Parameters:

cfg (KVConfig)
name (str)
joint_comm_group (JointGroups | None)
transfer_engine_info (TransferEngineInfo)
device (device)

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:

step (KVStep)
ctx (StepContext)

Return type:

AdmitOutcome

admit_retrieve(rid, node_name, graph_walk, published)[source]#

Takes the output of publish, possibly from another device, and kicks of a retrieval if needed (e.g., PD disaggregation KV transfer). Returns whether the retrieve has completed.

Parameters:

rid (str)
node_name (str)
graph_walk (str)
published (PublishedKVInfo | None)

Return type:

AdmitOutcome

classmethod build(spec, info)[source]#

Parameters:

spec (KVSpec)
info (EngineResourceInfo)

build_cuda_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:

slots (list[CGSlotSpec])
max_bs (int)
max_seq_len (int)

cleanup()[source]#

clear_preplan()[source]#

commit(step, ctx)[source]#

record step consumption

Parameters:

step (KVStep)
ctx (StepContext)

get_offload_priority(rid)[source]#

Device pages the request is holding — the most reclaimable first.

Parameters:: rid (str)
Return type:: float

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

Parameters:: overrides (KVReqConfig | None)

is_offloaded(rid)[source]#

True from the moment an offload claims the request, not just once its pages are on the host.

check_ready gates admission on this, so the window where the copy is still in flight must not look schedulable — and the worker’s victim filter must not pick a request that is already on its way out.

Parameters:: rid (str)
Return type:: bool

layer_view(layer_idx=None)[source]#

layer pages as needed by attention kernel

handed to AttentionManager::run. in kv_manager so storage mechanics are opaque to layers

Parameters:: layer_idx (int)
Return type:: Tensor

offload(rid)[source]#

Move every stream of rid to host memory. Returns pages freed.

A stream whose pages don’t fit on the host keeps them, so a partial offload still frees whatever did fit.

Device pages go back to the arena only once every stream has been copied: a step admitted before the claim can still run its fork copy, and that copy reads one of these streams.

Parameters:: rid (str)
Return type:: int

plan(step, ctx)[source]#

Returns list of sequence views per plan label

Parameters:

step (KVStep)
ctx (StepContext)

Return type:

dict[str, KVPlanOutput]

post_warmup_validate()[source]#

Assert num_free_pages is identical across every TP rank

Catches YAML drift (e.g. cpu_offload_pages set on one rank but not another), allocator-init bugs, and any future code path that adds requests asymmetrically before warmup returns. The all_gather itself is synchronizing, so no extra barrier is needed on the success path.

publish(request_id)[source]#

Parameters:: request_id (str)

read_kv(layer_idx=None, plan_label=None)[source]#

The slots this step’s plan writes, e.g. for NHD: [num_tokens, 2, num_kv_heads, head_dim] (K at index 0, V at 1).

Parameters:

layer_idx (int)
plan_label (str)

Return type:

Tensor

reclaimable(rid)[source]#

Device pages the request is holding; 0 once offloaded, and for one admitted but not yet run.

Parameters:: rid (str)
Return type:: int

reload(rid)[source]#

Bring every offloaded stream of rid back on device.

False when the device can’t fit them right now; nothing moves in that case, so the caller can evict further and try again.

Parameters:: rid (str)
Return type:: bool

remove_request(rid)[source]#

Parameters:: rid (str)

reset_default_cursors()[source]#

Return type:: None

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

For clearing dummy RIDs during cuda graph capture

Parameters:

rid (str)
free (bool)

set_layer_idx(layer_idx)#

Parameters:: layer_idx (int)
Return type:: None

property supports_eviction#

property supports_preplan#

write_kv(k, v, layer_idx=None, label=None, return_tensor=False)[source]#

Write K, V into this step’s planned slots.

Returns nothing by default: reading the slots back is a gather no caller wants today, and skipping it keeps the write a pure mutation.

Parameters:

k (Tensor)
v (Tensor)
layer_idx (int)
label (str)
return_tensor (bool)

Return type:

Tensor | None

class mstar.engine.resources.kv.manager.KVPlanState(token_to_page: torch.Tensor, token_to_cache: torch.Tensor, total_tokens: int | None = None)[source]#

Bases: object

Parameters:

token_to_page (Tensor)
token_to_cache (Tensor)
total_tokens (int | None)

copy_(other, capture_len)[source]#

Stage a step’s addressing into this captured state.

Neutralize only [n:capture_len] — the slots the graph scatters beyond the real tokens (SINK_PAGE, else they hit another request’s KV). Decode fills its bucket exactly (n == capture_len), so no-op there; only packed prefill pays it, over the real gap not the whole buffer.

Parameters:

other (KVPlanState)
capture_len (int)

token_to_cache: Tensor#

token_to_page: Tensor#

total_tokens: int | None = None#

class mstar.engine.resources.kv.manager.KVSequenceInfo(seq_len: int, latest_kv_transfer_info: Any, page_indices: list[int] = <factory>)[source]#

Bases: object

Parameters:

seq_len (int)
latest_kv_transfer_info (Any)
page_indices (list[int])

latest_kv_transfer_info: Any#

page_indices: list[int]#

seq_len: int#

class mstar.engine.resources.kv.manager.PageArena(kv_cache, allocator)[source]#

Bases: object

physical storage and free list management

Parameters:

kv_cache (KVCache)
allocator (PageAllocator)

acquire(n)[source]#

Parameters:: n (int)
Return type:: list[int] | None

allocator: PageAllocator#

copy_pages(src, dst)[source]#

Parameters:

src (list[int])
dst (list[int])

Return type:

None

kv_cache: KVCache#

property num_free#

release(pages)[source]#

Parameters:: pages (list[int])
Return type:: None

class mstar.engine.resources.kv.manager.PublishedKVInfo(info: dict[int, dict[str, mstar.engine.resources.kv.manager.KVSequenceInfo]]=<factory>, world_size: int = 1)[source]#