mstar.engine.resources.kv.transfer

mstar.engine.resources.kv.transfer#

KV transfer: moving one request’s pages between engines.

mstar.communication.tensors is imported lazily, inside the two places that genuinely need it at runtime (the transfer-engine isinstance dispatch and the Mooncake read). It pulls in the conductor and the sampling kernels behind it, and this module is otherwise free of both — keeping the import deferred is what lets the KV layer be built and tested without a GPU toolchain present.

Classes

CudaIpcKVTransferEngine(kv_cache[, max_workers])

CudaIpcKVTransferInfo(cuda_share, size, ...)

KVReadInfo(layer_idx, local_page_idx, ...)

KVTransferEngine()

KVTransferManager(transfer_engine_info, kv_cache)

LocalOnlyKVTransferEngine()

KV cache that remains local to its worker.

MooncakeKVTransferEngine(kv_cache, ...)

MooncakeKVTransferInfo(entity_id, ...)

TransferEngineInfo(my_entity_id, ...)

class mstar.engine.resources.kv.transfer.CudaIpcKVTransferEngine(kv_cache, max_workers=3)[source]#

Bases: KVTransferEngine

Parameters:

kv_cache (KVCache)

get_kv_transfer_info()[source]#
Return type:

CudaIpcKVTransferInfo

read_batched_async(remote_kv_info, read_info)[source]#
Parameters:
shutdown()[source]#
class mstar.engine.resources.kv.transfer.CudaIpcKVTransferInfo(cuda_share: 'tuple', size: 'tuple', stride: 'tuple', offset: 'int', dtype: 'str', requires_grad: 'bool', layout: 'KVLayout')[source]#

Bases: object

Parameters:
cuda_share: tuple#
dtype: str#
layout: KVLayout#
offset: int#
requires_grad: bool#
size: tuple#
stride: tuple#
class mstar.engine.resources.kv.transfer.KVReadInfo(layer_idx: 'int', local_page_idx: 'int', remote_page_idx: 'int', token_start: 'int', token_end: 'int')[source]#

Bases: object

Parameters:

  • layer_idx (int)

  • local_page_idx (int)

  • remote_page_idx (int)

  • token_start (int)

  • token_end (int)

layer_idx: int#
local_page_idx: int#
remote_page_idx: int#
token_end: int#
token_start: int#
class mstar.engine.resources.kv.transfer.KVTransferEngine[source]#

Bases: ABC

abstractmethod get_kv_transfer_info()[source]#
Return type:

Any

abstractmethod read_batched_async(remote_kv_info, read_info)[source]#
Parameters:

read_info (list[KVReadInfo])

Return type:

Future | None

abstractmethod shutdown()[source]#
class mstar.engine.resources.kv.transfer.KVTransferManager(transfer_engine_info, kv_cache)[source]#

Bases: object

Parameters:
cleanup()[source]#
get_kv_transfer_info()[source]#

Descriptor another process needs to read this cache remotely. KVCachePool.publish stamps it onto every SequenceInfo.

start_async_retrieve(start_len, end_len, local_page_indices, remote_page_indices, kv_transfer_info)[source]#
Parameters:

  • start_len (int)

  • end_len (int)

  • local_page_indices (list[int])

  • remote_page_indices (list[int])

  • kv_transfer_info (Any)

Return type:

Future | None

sync_retrieve(*args, **kwargs)[source]#
class mstar.engine.resources.kv.transfer.LocalOnlyKVTransferEngine[source]#

Bases: KVTransferEngine

KV cache that remains local to its worker.

get_kv_transfer_info()[source]#
Return type:

None

read_batched_async(remote_kv_info, read_info)[source]#
Parameters:

read_info (list[KVReadInfo])

Return type:

Future | None

shutdown()[source]#
class mstar.engine.resources.kv.transfer.MooncakeKVTransferEngine(kv_cache, entity_id, transfer_engine)[source]#

Bases: KVTransferEngine

Parameters:
get_kv_transfer_info()[source]#
Return type:

MooncakeKVTransferInfo

read_batched_async(remote_kv_info, read_info)[source]#
Parameters:
Return type:

Future | None

shutdown()[source]#
class mstar.engine.resources.kv.transfer.MooncakeKVTransferInfo(entity_id: 'str', session_id: 'str', data_ptr: 'int', layout: 'KVLayout')[source]#

Bases: object

Parameters:
data_ptr: int#
entity_id: str#
layout: KVLayout#
session_id: str#
class mstar.engine.resources.kv.transfer.TransferEngineInfo(my_entity_id: 'str', my_session_id: 'str', transfer_engine: 'TensorTransferEngine')[source]#

Bases: object

Parameters:
my_entity_id: str#
my_session_id: str#
transfer_engine: TensorTransferEngine#