Crate onnx_runtime_virtual_memory 

onnx-runtime-virtual-memory

Virtually contiguous, physically scattered memory.

The problem this solves

Paged KV storage and attention operators want opposite things. Paging wants small, individually reclaimable, individually migratable blocks. A GroupQueryAttention kernel wants one flat buffer per layer, because that is what the ONNX graph declares its past_key/past_value inputs to be.

The usual reconciliations are to copy the pages into a contiguous staging buffer every step, or to change the model graph so the operator understands block tables. The first costs a full KV copy per decode step; the second only works for models we control the export of.

There is a third option: reserve a contiguous range of virtual addresses and map physically separate blocks into it. The operator sees one flat buffer and runs unmodified. The blocks stay individually reclaimable, and growing the range costs a mapping call rather than a copy.

What decides whether this is cheap

Mapping granularity, which is a platform and device property rather than something this crate chooses. Measured on a Windows host with an RTX 4060:

mapping	granularity	tokens per granule, 8B GQA (2048 B/token)
Windows host	64 KiB	32
Linux / macOS host	page size, 4 KiB or 16 KiB	2 to 8
CUDA VMM	2 MiB	1024

On the host that is as fine as a KV page, so virtual contiguity costs nothing in wasted memory. On CUDA a sequence rounds up to num_layers * 2 * 2 MiB whatever its length, which only matters with many concurrent short sequences. See #596 for why that trade was accepted.

Apple Silicon

macOS uses the same mmap path as Linux. MAP_NORESERVE is defined there but is effectively ignored; it is an accounting hint, not a correctness requirement, so the reservation behaves the same. Note that Apple Silicon pages are 16 KiB, not 4 KiB, so granularity must be queried rather than assumed — a hard-coded 4096 would misalign every offset.

The bigger Apple consequence is not in this crate: CPU and GPU share one physical pool, so “device memory” and “host memory” are the same bytes. Anything holding separate per-tier budgets will over-commit there unless it knows the tiers alias.

What this crate is not

It does not decide whether memory may be held — that is onnx-runtime-memory-governor. A VirtualRange is a mapping mechanism; a lease is the permission to use one.

Re-exports

pub use backing::HostBacking;

pub use backing::PhysicalMemoryAccounting;

pub use backing::VirtualBacking;

pub use buffer::VirtualBuffer;

pub use buffer::VirtualBufferError;

Modules

backing

What a virtual address range is made of.

buffer

A buffer that grows without moving.

Structs

VirtualRange

A reserved range of virtual addresses with nothing behind it yet.

Enums

VirtualMemoryError

Why a virtual range could not be reserved, mapped, or unmapped.

Functions

granularity

This platform’s minimum mapping granularity, in bytes.