Unified Memory in SGLang: 1.9× as many prefixes retained, 37.9% less time for returning requests

Prerequisite: read Unified Radix Cache.

Background

During inference, a model saves keys and values, or KV, for use when generating the answer. The KV cache is the container that holds this KV. Modern hybrid models have different kinds of attention layers, and they save this history in different ways:

• Full-attention layers keep keys and values for each token. Reading more tokens adds more KV. During generation, attention reads this KV to look back at the document.
• Recurrent layers update a fixed-size state. Each new token is written into that state, overwriting its earlier contents: the contents change, but the state does not grow. Because of this overwriting, a state can be reused only if it was saved at exactly the same position; otherwise the history it held has already been overwritten. Each active request needs its own working state. Keeping a copy of that state for later reuse, called a checkpoint, takes another state slot.

The figure follows one request. As its context grows, KV accumulates while the recurrent state is updated in place. A hybrid model needs both to continue generating.

This creates a sizing problem. Many short requests need many state slots but little KV per request. A few long requests need fewer working states but much more KV. If SGLang reserves a fixed amount of memory for each kind at startup, one pool can fill up while the other still has unused space, so no fixed split handles both cases.

Unified Memory

Share the unused space

SGLang’s Unified Memory puts both kinds of data in one GPU buffer. Recurrent states grow from one end; KV grows from the other. Either side can take bytes from the free region between them. The total buffer size stays the same, while the capacity available to each kind changes with demand.

Suppose that buffer has 24 units. Each request needs 2 units of recurrent state and 1 unit of KV. A fixed split of 8 units for states and 16 for KV admits only four requests: the state pool is full, even though 12 KV units are unused. The shared pool fits all seven requests in the example, using 21 units.

Long contexts can cause the opposite problem: KV runs out while state capacity sits idle. Select “Long contexts” to see that case. Sharing helps when the fixed split is the limit. When the combined data fills the whole buffer, requests still have to wait or cached data must be evicted.

Here, “Unified Memory” means sharing capacity between SGLang’s caches on the same GPU. It does not move data to CPU memory. HiCache handles that separate task by saving cached history in CPU memory or external storage and loading it back when needed.

Pack data into movable blocks

The shared buffer is one continuous range of GPU memory. The allocator hands out space from its two ends. It also has to return freed space to the free region in the middle, and as shown below, it sometimes moves blocks to close holes. All of this requires each block it manages to occupy one continuous range.

KV is not stored that way by default. SGLang allocates KV in pages, each holding several tokens, and every full-attention layer keeps its own KV for those tokens. Normally each layer has a separate buffer, so one page is split into one piece per layer, spread across those buffers. Releasing the page leaves a small hole in every layer’s buffer, so the freed space stays scattered across the separate layer buffers.

Unified Memory therefore switches to a page-major layout. All layers’ KV for one page sits side by side in one contiguous block, called a page envelope. The data in the page is unchanged; only its position in memory differs. Releasing the page now frees one continuous range, which can return to the shared free region. Moving the page copies one block instead of one piece per layer.

This different layout has costs:

• Kernels must read each layer through strided views into the envelopes, so Unified Memory currently supports only some attention backends and requires K and V heads of the same size.
• Each forward pass must first translate virtual IDs into physical positions.

The measured overhead is covered under Results at the end.

The figure compares the two layouts using three layers and two pages.

A recurrent-state block similarly holds the state needed across the recurrent layers. KV pages and state blocks keep their own sizes. Both draw from the same buffer, so freeing a large state block can make room for several smaller KV pages.

Consolidate freed space at runtime

In Unified Memory, a freed block may sit between two live blocks. In the figure, state blocks A, B and C occupy the left end. Releasing B leaves a hole, but C still separates that hole from the shared free region, so the freed space cannot be fully used. SGLang therefore moves C into B’s old location. C’s old space then joins the free region, where new KV pages can be allocated.

Requests still need to find C after it moves. SGLang gives C a stable virtual ID and keeps a table of its current location. The request continues to refer to C; the table changes from “C is at 8” to “C is at 4”. Moving the data does not require rewriting every request that refers to it.

This process is called compaction. SGLang can leave a hole temporarily and compact when an allocation needs the space. Moving data must wait until it is safe: data still in use by GPU computation or a network transfer cannot be moved in the middle of that operation.

The shared-pool implementation introduced this allocation and mapping mechanism. A later extension lets three kinds share the buffer: full-attention KV, sliding-window KV and recurrent state. Sliding-window layers only need a recent window of tokens; their capacity demand can differ from both of the other pools.

Eviction policy in Unified Memory

After a request finishes, SGLang can keep the document’s KV and a recurrent-state checkpoint from the end of the request. When another request starts with exactly the same token prefix, it can reuse that saved history and compute only the new suffix. When memory is limited, evicting the cached prefix frees space immediately, but a later request may have to compute it again.

With shared memory, the eviction decision must count all the capacity it can reclaim. Suppose a new allocation needs 6 units of KV space. Evicting one cached prefix releases 1 unit of KV and a 2-unit state checkpoint. After two evictions and any needed compaction, (1 + 2) × 2 = 6 units are available. Counting only the 2 KV units would make SGLang keep deleting prefixes even though the allocation already fits.

SGLang’s eviction improvement checks whether the shared allocator can satisfy the pending allocation after each eviction. Once it can, eviction stops. The check includes space recoverable from the other pool, so more useful prefixes can remain cached.

Results

The experiment below first caches 28 different prefixes, then sends an unrelated long request to create memory pressure. Finally, it sends the original prefixes again. A surviving prefix saves repeated prompt computation; an evicted one must be recomputed.

Both variants already use Unified Memory. The comparison changes when allocation-driven eviction stops. With the improved rule, more prefixes survive, and the 28 returning requests finish sooner.

This gain has a precondition: a single allocation must be large. In the experiment, the 7,000-token request allocates all its KV at once. When it is split into four 2,048-token chunks, the old and new rules give the same result, both keeping 17 prefixes. Overall throughput showed no stable change.

Unified Memory also has costs: the location table, block copies during compaction and the layout costs described earlier. Across five models, median time between output tokens ranged from 0.81% faster to 1.32% slower. The benefit is better use of memory, not faster computation.