Reserve KV Cache for the Output You Promised, Not Just the Prompt You Received

A developer detailed a fix for out-of-memory errors in self-hosted LLM servers, arguing that admission control must reserve KV cache based on the worst-case generation cap (prompt plus max output tokens) rather than prompt length alone. The post provides code for calculating per-token KV memory and clamping unspecified max_tokens to the model's context limit, and recommends returning Retry-After on rejection to prevent retry storms.

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.