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:
def kv_bytes_per_token(num_layers, num_kv_heads, head_dim, bytes_per_element=2):
return 2 * num_layers * num_kv_heads * head_dim * bytes_per_element
def request_reservation_tokens(prompt_tokens, max_output_tokens):
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
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.
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.
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.