onnx-runtime-virtual-memory 0.1.0-dev.6

Virtually contiguous, physically scattered memory: reserve one address range and map separately-allocated blocks into it, so an operator that requires a flat buffer can read paged storage without a copy
Documentation

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.