--swa-full-tokens-ratio, --mamba-full-memory-ratio). When the workload does not match that split, one pool runs out while the other has room to spare.
With --enable-unified-memory, every kind lives in one device buffer, the arena. Each kind is a sub-pool that occupies a contiguous band of the arena’s bytes, and bands grow and shrink as requests come and go. Memory moves between kinds as the workload shifts: long prompts take more full-attention KV, many short concurrent requests take more state slots and window KV.
This page covers when and how to enable it. The design page describes how it works and how to debug and test it.
What changes
Sub-pools and bands
Each kind of memory a model’s layers declare becomes one sub-pool:
Sub-pools are laid out in the arena by retention: the state sub-pool takes the low end and grows up, the full-attention sub-pool takes the high end and grows down, and window sub-pools sit between them. With no state sub-pool, the longest window takes the low end. A sub-pool between the two ends is a float: it occupies a span between its neighbours and is relocated when the layout needs it.
Supported models
Unified memory replaces the pools of hybrid models only. At startup,KVCacheConfigurator._kv_layer_decls declares the model’s layers and refuses anything else.
The target model’s sliding-window layers form one window sub-pool. A fused draft model whose window has a different length adds a second one. Only the Triton attention backend serves several window sub-pools per runner; another backend bound to more windows than it supports raises at startup, naming itself.
Attention, linear-attention and Mamba backends
Page-major entries make each layer’s rows strided, so only backends audited for strided views are accepted. The table is that page-major allow-list (arg_groups/kv_cache_hook.py::handle_page_major_kv_layout); a backend on it may still refuse a model when it is built:
fa3 and fa4 still refuse two kinds of model when the backend is built: local-attention models, such as Llama 4, whose default backend on SM90 is fa3 (FlashAttentionBackend._maybe_init_local_attn_builder), and models that decode with prefill-aware SWA (FlashAttentionBackend.__init__).
If the FULL prefill CUDA-graph backend is configured with an attention backend other than fa3 or fa4, it is disabled with a warning; decode capture and the other prefill graph backends are unaffected.
Quick start
Platform: NVIDIA CUDA. On other platforms, read Known limitations first.MODEL_PATH is a Hugging Face repo id or local path of a supported hybrid model.
Flags and environment variables
A prompt that cannot fit even an otherwise empty arena is aborted at admission with
Request prompt exceeds the KV memory budget. The memory clip and admission price a request with the same demand model (ChainAllocator.max_new_tokens_for_memory). On a model with sliding-window and Mamba layers whose cache keeps state checkpoints, its window bound assumes the checkpoints advance; when no tracking or replacement state slot is free, a request can hold more window rows than it was priced for, so passing the clip does not guarantee that its chunks finish (see Known limitations).
Interaction with other features
Refused at startup:
--enable-two-batch-overlap--enable-lmcache,--enable-unified-cache-external-linker--enable-kv-cache-sharding--enable-int8-mamba-checkpoint--enable-prefill-weight-versions- FP4 KV cache (
--kv-cache-dtype nvfp4orfp4_mx_block16) --cp-kv-layoutvalues that keep each rank’s KV sharded and padded- out-of-tree platforms
Observability
Pool usage is reported per sub-pool: each window sub-pool has its own label, and state is reported in slots. Because sub-pools share bytes, the available figures are estimates; the held figures are exact. At idle, the scheduler checks every sub-pool for leaked pages and the arena for byte conservation. The design page lists the stricter checks and the test tools.Known limitations
- Priority preemption. With
--enable-priority-scheduling, unified memory turns priority preemption off, with a startup warning: a higher-priority request is admitted first but does not retract running requests, so under load it waits for running requests to finish. - MPS and MLX. Refused at startup: unified memory orders the scheduler’s writes after the running forward’s reads with device stream waits, and the scheduler has no device streams there.
- Window rows when state checkpoints stall. On a model with sliding-window and Mamba layers whose cache keeps state checkpoints, the memory clip and admission assume each request’s checkpoints advance. When no tracking or replacement state slot is free, a checkpoint stalls and the request holds more window rows than it was priced for, so an admitted request is not guaranteed to finish. See the checkpoint retention bound.
- PD failed transfers. A failed or timed-out transfer’s pages are released at its deferred-release timeout, or a failed send’s at once, without proof that the remote transfer engine stopped touching them.