# Reserve KV Cache for the Output You Promised, Not Just the Prompt You Received > Source: > Published: 2026-08-12 03:08:53+00:00 If your self-hosted LLM server checks admission against the prompt length alone, it will still OOM — because a short prompt with a large `max_tokens` cap is a bigger memory promise than a long prompt with a tight one. The fix is to admit requests against the **worst case they are allowed to reach** (prompt + full generation budget), reserve that memory up front, and when you have to reject, return `Retry-After` so clients back off instead of hammering you. This post is the accounting that makes that work. I'm writing this as a deeper follow-up to an earlier post on admission control for LLM serving. A commenter there made a point sharp enough to deserve its own walkthrough: the generation reserve is the part that saves you from the weird, intermittent failures, and admission control without a backpressure signal just turns into a retry storm with nicer logs. Both are right. Here's how to build them. The KV cache — the cached key/value tensors for every token in a request's context — grows with sequence length. The catch is that "sequence length" is not the prompt you received; it's the prompt **plus every token you're still going to generate**. A request that arrives with 200 prompt tokens and `max_tokens=4000` will, if the model runs to that cap, occupy KV cache for 4,200 tokens before it finishes. Admit it based on the 200 you can see and you've under-counted its footprint by 20×. This is why the failures are intermittent and infuriating. Most requests stop generating early, so the optimistic accounting appears to work for days. Then several long-generating requests overlap, each quietly climbing toward its cap, and the cache OOMs — not at admission, but somewhere mid-decode, taking down in-flight requests that were already accepted. The crash has no obvious trigger because the request that *caused* it looked tiny when you let it in. The mental model that fixes this: **admission is a promise about the future, so it must be priced at the maximum the request is allowed to cost, not the minimum it costs right now.** **Takeaway:** A request's memory footprint is bounded by its generation cap, not its prompt — admit on the cap or you're admitting a bill you haven't read. Reserve for the full context the request can reach: prompt tokens plus its output cap. The per-token KV cost is fixed by the model architecture: ``` python def kv_bytes_per_token(num_layers, num_kv_heads, head_dim, bytes_per_element=2): # 2 = keys + values. num_kv_heads, not total heads: # grouped-query attention (GQA) shares KV across query heads, # shrinking this 4-8x vs the naive per-head formula. return 2 * num_layers * num_kv_heads * head_dim * bytes_per_element def request_reservation_tokens(prompt_tokens, max_output_tokens): # The worst case this request is allowed to reach. return prompt_tokens + max_output_tokens def request_reservation_bytes(prompt_tokens, max_output_tokens, per_token): return request_reservation_tokens(prompt_tokens, max_output_tokens) * per_token ``` The number that matters is `max_output_tokens` , and the trap is that clients love to leave it unset. An unset cap is not "small" — it's the model's context-window maximum, which is the largest promise a single request can make. Any admission controller that treats a missing cap as zero is guaranteeing itself an OOM. Clamp it: ``` MODEL_CONTEXT = 8192 # the model's max context, as of your deployment def effective_output_cap(requested_max_output, prompt_tokens): remaining = MODEL_CONTEXT - prompt_tokens if remaining <= 0: return None # prompt alone exceeds context: reject, don't truncate silently # An unset cap means "up to the context limit", not "small". requested = requested_max_output if requested_max_output is not None else remaining return min(requested, remaining) ``` **Takeaway:** Treat an unspecified `max_tokens` as the context maximum, because that's exactly what it can become in production. Admission control is a running ledger, not a one-time check. You subtract a request's reservation from the free budget when you accept it and add it back the moment it completes, fails, or is cancelled. The reservation is held for the request's *entire* lifetime at its worst-case size — you don't shrink it as tokens generate, because the whole point is to guarantee the request can reach its cap without a mid-flight OOM. ``` python import threading class KVAdmissionController: def __init__(self, total_kv_bytes, per_token_bytes, safety_fraction=0.9): self.usable = int(total_kv_bytes * safety_fraction) # leave headroom self.per_token = per_token_bytes self.reserved = 0 self.lock = threading.Lock() def try_admit(self, prompt_tokens, output_cap): need = (prompt_tokens + output_cap) * self.per_token with self.lock: if self.reserved + need <= self.usable: self.reserved += need return need # a handle to release later return None # rejected: no budget def release(self, need): with self.lock: self.reserved = max(0, self.reserved - need) ``` Two details that bite people. First, `release` must run in a `finally` block — if a request errors out and you skip the release, that budget leaks permanently and your server slowly strangles itself until a restart. Second, the `safety_fraction` is not optional padding; real serving runtimes carry fragmentation and activation overhead beyond the KV cache, so admitting to 100% of theoretical capacity will OOM at maybe 92% of it. I keep the usable budget around 85–90% and tune down if I still see pressure. **Takeaway:** Hold the reservation at full size for the whole request and release it in `finally` — a leaked reservation is a slow-motion outage. `Retry-After` matter so much? A rejection without guidance is an invitation to retry immediately, and a client library retrying immediately against a saturated server produces a retry storm: the requests most likely to be rejected are the ones that come back fastest, so load *concentrates* exactly when you have the least to spare. The HTTP-native fix is a `429 Too Many Requests` with a `Retry-After` header telling the client how long to wait. ``` python from fastapi import FastAPI, Request from fastapi.responses import JSONResponse app = FastAPI() @app.post("/generate") async def generate(req: Request): body = await req.json() prompt_tokens = count_tokens(body["prompt"]) cap = effective_output_cap(body.get("max_tokens"), prompt_tokens) if cap is None: return JSONResponse({"error": "prompt exceeds context window"}, status_code=413) handle = controller.try_admit(prompt_tokens, cap) if handle is None: return JSONResponse( {"error": "server at capacity", "retry_after_seconds": 2}, status_code=429, headers={"Retry-After": "2"}, ) try: return await run_generation(body, prompt_tokens, cap) finally: controller.release(handle) ``` `Retry-After` works because well-behaved HTTP clients — including most SDK retry layers — honor it, converting a thundering herd into staggered arrivals. Pair it with jittered exponential backoff on the client and the retries spread out instead of synchronizing. The value can be a fixed conservative number or an estimate derived from your average request duration; even a static `2` beats no signal at all. Whatever you choose, the header is the contract: it turns "I said no" into "come back at a time when I might say yes." **Takeaway:** A rejection without `Retry-After` doesn't shed load, it reschedules it for one second from now — the header is what actually protects you. | Strategy | What it counts | Fails when | Backpressure | |---|---|---|---| | Concurrency cap (max N requests) | Request count only | Long-context requests, mixed sizes | Usually none | | Prompt-length admission | Prompt tokens | Short prompt, large `max_tokens` | Often none | | Prompt + output reservation | Prompt + generation cap | Mostly holds; over-reserves early stops | `429` + `Retry-After` | A fixed concurrency cap is the common starting point and the weakest: eight requests can mean 8K tokens or 80K depending on their caps. Prompt-length admission is better but still blind to the generation promise. Reserving for prompt-plus-output is the one that survives adversarial traffic; its only real cost is that requests which stop early briefly over-reserve, which you recover the instant they finish. **Should I reserve KV cache for max_tokens or the actual output length?** Reserve for `max_tokens` (the cap), because you cannot know the actual output length at admission time and the request is allowed to run all the way to that cap. Reserving for a hoped-for shorter length is exactly the optimistic accounting that OOMs under load. **What should an LLM server return when it's out of KV cache?** Return HTTP `429 Too Many Requests` with a `Retry-After` header. The status code tells the client it's a transient capacity issue worth retrying, and the header staggers those retries so they don't synchronize into a storm against an already-saturated server. **Why does a short prompt cause more trouble than a long one?** Because footprint is prompt length plus the generation cap, and short prompts often carry large or unset generation caps. A 200-token prompt with `max_tokens=4000` reserves for 4,200 tokens; a 3,000-token prompt with `max_tokens=200` reserves for 3,200. The short one is the bigger promise. If you self-host an LLM behind real traffic, admit requests against their worst case — prompt tokens plus the output cap, with an unset cap treated as the context maximum — and hold that reservation for the request's whole life, releasing it in a `finally` block. When there's no budget, reject with `429` and a `Retry-After` header so clients back off instead of stampeding. Start the usable budget near 85–90% of theoretical KV capacity to absorb fragmentation, and tune from there. The concurrency cap you probably started with is fine as a coarse safety net, but it's the output reservation and the backpressure signal that keep the server up when the traffic gets weird.