mstar.engine.resources.kv.cpu_page_pool

mstar.engine.resources.kv.cpu_page_pool#

Host-side page pool for KV cache offloading.

When device pages run out, a cold request’s pages move here, into pinned host memory, freeing device pages for the requests that are actually running. They move back when it is scheduled again.

The pool stores pages and the stream metadata that has to survive the round trip; which request to evict, and when, is the worker’s decision.

Classes

CPUPagePool(config, kv_cache, max_cpu_pages)

Host mirror of the paged KV cache, holding whatever is offloaded.

OffloadedStream(cpu_page_indices, ...[, ...])

One (request, label) stream living on the host.
class mstar.engine.resources.kv.cpu_page_pool.CPUPagePool(config, kv_cache, max_cpu_pages)[source]#

Bases: object

Host mirror of the paged KV cache, holding whatever is offloaded.

Parameters:
discard(rid, label)[source]#

Drop a stream that was copied here but never committed to.

The offload aborted after the copy, so the device pages stay live and these host pages go back to the pool.

Parameters:
Return type:

None

is_offloaded(rid)[source]#
Parameters:

rid (str)

Return type:

bool

labels(rid)[source]#
Parameters:

rid (str)

Return type:

list[str]

property num_free_pages: int#
num_pages(rid, label)[source]#
Parameters:
Return type:

int

offload_stream(rid, label, gpu_kv_cache, gpu_page_indices, stored_len, position, released=0)[source]#

Copy device pages to the host. False when the host pool is full, in which case nothing moved and the caller keeps its device pages.

Parameters:
Return type:

bool

reload_stream(rid, label, gpu_kv_cache, gpu_page_indices)[source]#

Copy the host pages back onto gpu_page_indices and release them.

Returns the stream metadata saved at offload, for the caller to restore.

Parameters:
Return type:

OffloadedStream

remove_request(rid)[source]#
Parameters:

rid (str)

Return type:

None

sync()[source]#

Order the current stream behind the pending copies.

Both directions need this before the pages they touched are reused: after an offload the device pages go back to the allocator, and after a reload the attention kernels read them.

Return type:

None

class mstar.engine.resources.kv.cpu_page_pool.OffloadedStream(cpu_page_indices, stored_len, position, released=0)[source]#

Bases: object

One (request, label) stream living on the host.

Parameters:
cpu_page_indices: list[int]#
position: int#
released: int = 0#
stored_len: int#