mstar.engine.resources.kv.cache

Classes

KVCache(cfg, device[, dtype])	Owns the KV storage and all layout-dependent addressing, so that consumers (KV transfer, attention) stay layout-agnostic.
PageAllocator(max_num_pages)	Simple page allocator using a FIFO queue of free page indices.

class mstar.engine.resources.kv.cache.KVCache(cfg, device, dtype=torch.bfloat16)

Bases: object

Owns the KV storage and all layout-dependent addressing, so that consumers (KV transfer, attention) stay layout-agnostic.

Parameters:

• cfg (KVConfig)

• device (device)

chunk_ptrs(layer_idx, page_idx, token_start, token_end, base_ptr=None)

Byte pointers to the contiguous chunks holding tokens [token_start, token_end) of one page of one layer, plus the size of each chunk. One pointer per chunk (K and V are separate under NHD).

base_ptr addresses a remote cache with this same layout/config; it defaults to this cache’s own storage.

Parameters:

• layer_idx (int)

• page_idx (int)

• token_start (int)

• token_end (int)

• base_ptr (int | None)

Return type:

tuple[list[int], int]

chunk_view(layer_idx, page_idx, token_start, token_end, tensor=None)

View of tokens [token_start, token_end) of one page of one layer, covering both K and V. tensor overrides the backing storage (e.g. a tensor rebuilt from another process’ cache with this same layout).

Parameters:

• layer_idx (int)

• page_idx (int)

• token_start (int)

• token_end (int)

• tensor (Tensor | None)

Return type:

Tensor

copy_pages(src_pages, dst_pages)

Copy whole pages (every layer, both K and V, all tokens) within this cache: src_pages[i] -> dst_pages[i].

Parameters:

• src_pages (list[int])

• dst_pages (list[int])

Return type:

None

data_ptr()

Return type:

int

property device: device

property dtype: dtype

layer_view(layer_idx)

One layer’s pages, in this cache’s layout — what an attention kernel consumes. NHD: [max_num_pages, 2, page_size, num_kv_heads, head_dim].

Parameters:

layer_idx (int)

Return type:

Tensor

property layout: KVLayout

property nbytes: int

property num_layers: int

property page_size: int

read_tokens(layer_idx, page_idx, cache_idx)

Gather the (page, offset-in-page) slots written by write_tokens. Returns [num_tokens, 2, num_kv_heads, head_dim] (K at index 0, V at 1); it is a gather, so a copy rather than a view.

Parameters:

• layer_idx (int)

• page_idx (Tensor)

• cache_idx (Tensor)

Return type:

Tensor

write_tokens(layer_idx, k, v, page_idx, cache_idx, return_tensor=False)

Scatter per-token K/V ([num_tokens, num_kv_heads, head_dim]) into the (page, offset-in-page) slots given by page_idx/cache_idx.

Parameters:

• layer_idx (int)

• k (Tensor)

• v (Tensor)

• page_idx (Tensor)

• cache_idx (Tensor)

• return_tensor (bool)

Return type:

None

class mstar.engine.resources.kv.cache.PageAllocator(max_num_pages)

Bases: object

Simple page allocator using a FIFO queue of free page indices.

Thread-safe: a threading.Lock makes the qsize-then-get sequence in allocate/try_allocate atomic against concurrent free calls. Required by the pre-plan path, where the plan thread runs try_allocate while the GPU thread runs free from reset_label — the unlocked qsize/get pair could false-negative (return None when pages are about to be freed) or partially fill the output list under multi-consumer contention.

Parameters:

max_num_pages (int)

allocate(n)

Parameters:

n (int)

Return type:

list[int]

free(pages)

Parameters:

pages (list[int])

Return type:

None

property num_free: int

try_allocate(n)

Like allocate() but returns None instead of raising on failure.

Parameters:

n (int)

Return type:

list[int] | None