mstar.engine.resources.kv.page_copier

Whole-page copies between the device KV cache and the host page pool.

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 page 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 CudaBindingsCopier, then CudartCopier; 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, for torch 2.9, which ships no cuda.bindings.
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, for torch 2.9, which ships no cuda.bindings.

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

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)

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