Inference Systems · Accelerator Scheduling

Admission Control and the Memory Watermark

Reserving headroom so that growth does not become preemption: how watermarks, reservation bounds, and head-of-line blocking interact in a cache pool.

Every running sequence will need more memory than it holds. A decode sequence appends one token’s keys and values per step, so its residency grows monotonically until it stops. Admission control is the discipline of deciding, at the moment a request arrives, whether the pool can afford to carry it to completion. The decision is made with an estimate, and the quality of the estimate determines whether the system spends its time generating tokens or recovering from memory pressure.

Why admission cannot wait for pressure

The tempting alternative to admission control is to admit everything and handle exhaustion when it happens. This fails for a structural reason: recovering memory from a running sequence is expensive, and the cost falls on every other sequence in the batch.

Recovering by swapping a sequence’s cache to host memory means copying the cache out and later back. At 128 KiB per token, a 2,000-token sequence is 256 MiB, and moving it over a PCIe link takes milliseconds even at full link rate. During that move the device is not generating anything. Recovering by discarding the cache and recomputing it means re-running the prompt and any generated prefix through the model, which costs a full prefill of the sequence’s length. Both are order-of-magnitude more expensive than the allocation decision that would have prevented the situation.

An admission decision is cheap: a few arithmetic operations against a counter. The asymmetry is the entire argument. Admission control is a way of paying a small cost at the right time to avoid a large cost at the wrong time.

The watermark

The standard mechanism is to keep a reserve. Let the pool hold $M$ blocks, and let the scheduler admit a request only while the number of free blocks remains above a watermark $W$. The watermark serves two purposes at once: it absorbs the growth of already-admitted sequences between scheduling decisions, and it gives the allocator a buffer in which to satisfy the next block request without preempting anything.

Sizing $W$ requires an estimate of the maximum growth rate across the running set. If the scheduler makes an admission decision once per step, and each running sequence consumes one token’s worth of cache per step, then between decisions the pool loses $Bc$ bytes, where $B$ is the running batch and $c$ the per-token cost. The watermark must therefore exceed $Bc$ plus the largest single allocation a step can demand. For a batch of 128 sequences at 128 KiB per token, one step’s growth is 16 MiB, so a watermark below that figure is not a reserve at all.

A second component of the watermark covers the variance in what arrives. A request whose prompt is longer than the current chunk budget will be processed over several steps, and its cache grows with each chunk. If the scheduler admits several such requests at once, the growth between decisions is a multiple of the single-request estimate.

Admission bounds and their honesty

Three bounds are available for deciding whether a request fits, and they differ in how much they assume about the future.

Declared maximum. Reserve $\lceil T_{max}/P \rceil$ blocks at admission and never admit beyond the pool. This is exact and conservative: no request can ever fail after admission, because its full allowance is already accounted for. It is also pessimistic, because most requests stop well before their maximum, and the pool sits half empty in block terms while the queue grows. For a workload where the declared maximum is a large default and actual lengths are short, this bound can cut throughput by a large factor.

Current need only. Reserve only the blocks the prompt requires and allocate as the sequence grows. This maximises the number of sequences admitted and guarantees that admission never over-reserves. It provides no guarantee of completion: a sequence may run out of memory at any point, and the failure rate depends on the mix of lengths in the running set. This is the policy that produces preemption under load.

Growth allowance. Reserve the prompt’s blocks plus an allowance of $A$ tokens’ worth of growth, and extend the reservation when the sequence passes it. The allowance is chosen from the length distribution of the workload, so that most sequences finish within it. This is the middle ground and the one most systems settle on, because it caps the number of extension events, and each extension is a decision point where the scheduler might have to evict.

The honest framing is that only the first bound is a guarantee. The other two are heuristics that work when the length distribution is stable and degrade when it shifts. A system that switches from short chat responses to long document generations changes the distribution without changing any configuration, and the reservation policy that was adequate becomes the source of a preemption storm.

Head-of-line blocking in the waiting queue

Admission control interacts badly with first-come-first-served ordering. If the head of the waiting queue is a request that needs 20,000 tokens of cache and only 15,000 tokens’ worth is free, a strict policy holds every subsequent request behind it even if those requests are short and would fit.

The standard remedy is to scan past the head. Admitting a smaller request from further back in the queue uses memory that would otherwise be idle and increases the number of sequences making progress. The cost is a fairness violation: the large request waits while later arrivals proceed. Whether that is acceptable depends on whether the queue is ordered by arrival or by some notion of priority, and on whether the large request has a deadline.

A stricter variant preserves order and instead reserves the large request’s full need, holding the free memory idle until enough accumulates. This bounds the latency of the large request at the cost of throughput during the wait. Both policies are defensible; the failure is to have no policy and let the head block indefinitely while the pool is used by nothing.

The interaction with prefill chunks

Chunked prefill makes admission harder, because a request’s cache must be reserved for its full prompt before the first chunk runs, even though the blocks are only written as chunks complete. Admitting on the first chunk’s need and extending later means the request can fail midway through prompt processing, after having consumed device time.

The failure is worse than an admission rejection because the consumed time is unrecoverable and the request must be restarted from the beginning. A scheduler that chunks prompts therefore reserves against the full prompt length plus a growth allowance, which makes prompt-heavy workloads more conservative in admission than decode-heavy ones for the same pool. This is one of the reasons prefill and decode are often given separate pools; splitting prefill from decode inside a step covers the scheduling side.

A reservation check

The decision itself is small. The parts that matter are which bound is used and what happens on failure.

def try_admit(req, pool, watermark):
    blocks = ceil_div(req.prompt_tokens + req.growth_allowance,
                      pool.block_tokens)
    if pool.free < blocks + watermark:
        return False                      # do not overcommit the reserve
    for _ in range(blocks):
        pool.take()                       # marks blocks reserved, not filled
    req.reserved = blocks
    return True

def extend(seq, pool, watermark):
    if seq.tokens_written % pool.block_tokens != 0:
        return True                       # still inside the last block
    if pool.free <= watermark:
        return False                      # caller must preempt or drop
    pool.take()
    seq.reserved += 1
    return True

Two details in that code carry most of the semantics. The watermark is subtracted from free blocks rather than compared against them, so the reserve is never spent on admissions. And extend failing is not an error path; it is a normal outcome that the caller must handle, which means every decode step has to be prepared for a sequence that cannot advance.

Measuring whether the policy works

The observable quantities are the number of admissions rejected while free memory existed above the watermark, the number of sequences preempted, and the fraction of the pool held as reserve. A policy that rejects admissions with substantial free memory is too conservative; a policy that preempts frequently is too aggressive; a policy with a large reserve is paying throughput for a guarantee it may not need.

The distribution matters more than the mean. Admission failures cluster at the moment the longest-running sequences are at their peak, which is also the moment the queue is deepest. A policy tuned to average conditions will fail exactly when the system is busiest, which is when the failure is most visible. Sizing the watermark and the growth allowance for the tail of the length distribution, rather than the mean, is the difference between a system that degrades gracefully and one that collapses under a burst.

The last resort when admission and reservation both fail is to take memory back from a sequence that is already running. That is preemption, and the choice of how to take it back is a separate mechanism with its own cost structure, covered in preemption by swap versus recomputation.

References

  1. Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve — Agrawal et al., OSDI 2024
  2. Site Reliability Engineering: Chapter on Handling Overload — Beyer et al., O'Reilly 2016
  3. Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al., SOSP 2023