mstar.engine.resources.kv.page_copier

Page copies between KV caches: the device cache and the host page pool, or two device caches (CUDA IPC KV transfer).

Both caches are laid out [layers, pages, ...], so a page is one contiguous block per layer, a layer stride apart. copy_ of that strided view between device and pinned host memory blocks the host until the copy stream drains and runs at a fraction of the link’s bandwidth. On CUDA and ROCm each run of pages adjacent in both caches is instead one cudaMemcpy2DAsync / hipMemcpy2DAsync with the layer stride as its pitch: async, and run on the copy engines.

make_page_copier picks the first copier that loads and passes a round trip, once per process: on CUDA CudartCopier, then CudaBindingsCopier, whose calls drop the GIL; on ROCm HipCopier; then TorchCopier.

Functions

make_page_copier()

The fastest copier that works here; TorchCopier off CUDA and ROCm.

Classes

CudaBindingsCopier()	Through cuda.bindings, which Linux CUDA builds of torch 2.10+ require.
CudartCopier()	libcudart, the runtime torch loaded, through ctypes.PyDLL so each call keeps the GIL.
HipCopier()	libamdhip64 on ROCm, where copy_ of a strided page blocks the host too.
PageCopier()	Copies every layer of src_pages[i] of src onto dst_pages[i] of dst, queued on stream.
TorchCopier()	Per-page copy_: works everywhere, but on CUDA blocks the host.

class mstar.engine.resources.kv.page_copier.CudaBindingsCopier

Bases: _Memcpy2DCopier

Through cuda.bindings, which Linux CUDA builds of torch 2.10+ require.

name = 'cuda.bindings'

class mstar.engine.resources.kv.page_copier.CudartCopier

Bases: _MappedRuntimeCopier

libcudart, the runtime torch loaded, through ctypes.PyDLL so each call keeps the GIL. Dropping it costs up to a switch interval per call while another thread runs Python. A call that waits on the driver stalls every Python thread.

name = 'libcudart'

class mstar.engine.resources.kv.page_copier.HipCopier

Bases: _MappedRuntimeCopier

libamdhip64 on ROCm, where copy_ of a strided page blocks the host too.

name = 'libamdhip64'

class mstar.engine.resources.kv.page_copier.PageCopier

Bases: object

Copies every layer of src_pages[i] of src onto dst_pages[i] of dst, queued on stream.

copy(dst, dst_pages, src, src_pages, stream)

Parameters:

• dst (Tensor)

• dst_pages (list[int])

• src (Tensor)

• src_pages (list[int])

• stream (Stream)

Return type:

None

copy_rows(dst, src, stream)

dst.copy_(src) queued on stream, for two views of the same shape whose rows ([i]) are each contiguous: e.g. a page or a token range of it, every layer.

Parameters:

• dst (Tensor)

• src (Tensor)

• stream (Stream)

Return type:

None

name = 'base'

class mstar.engine.resources.kv.page_copier.TorchCopier

Bases: PageCopier

Per-page copy_: works everywhere, but on CUDA blocks the host.

copy(dst, dst_pages, src, src_pages, stream)

copy_rows(dst, src, stream)

dst.copy_(src) queued on stream, for two views of the same shape whose rows ([i]) are each contiguous: e.g. a page or a token range of it, every layer.

name = 'torch'

mstar.engine.resources.kv.page_copier.make_page_copier()

The fastest copier that works here; TorchCopier off CUDA and ROCm.

Return type:

PageCopier