{"slug": "reserve-kv-cache-for-the-output-you-promised-not-just-the-prompt-you-received", "title": "Reserve KV Cache for the Output You Promised, Not Just the Prompt You Received", "summary": "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.", "body_md": "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`\n\ncap 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`\n\nso clients back off instead of hammering you. This post is the accounting that makes that work.\n\nI'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.\n\nThe 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`\n\nwill, 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×.\n\nThis 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.\n\nThe 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.**\n\n**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.\n\nReserve 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:\n\n``` python\ndef kv_bytes_per_token(num_layers, num_kv_heads, head_dim, bytes_per_element=2):\n    # 2 = keys + values. num_kv_heads, not total heads:\n    # grouped-query attention (GQA) shares KV across query heads,\n    # shrinking this 4-8x vs the naive per-head formula.\n    return 2 * num_layers * num_kv_heads * head_dim * bytes_per_element\n\ndef request_reservation_tokens(prompt_tokens, max_output_tokens):\n    # The worst case this request is allowed to reach.\n    return prompt_tokens + max_output_tokens\n\ndef request_reservation_bytes(prompt_tokens, max_output_tokens, per_token):\n    return request_reservation_tokens(prompt_tokens, max_output_tokens) * per_token\n```\n\nThe number that matters is `max_output_tokens`\n\n, 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:\n\n```\nMODEL_CONTEXT = 8192  # the model's max context, as of your deployment\n\ndef effective_output_cap(requested_max_output, prompt_tokens):\n    remaining = MODEL_CONTEXT - prompt_tokens\n    if remaining <= 0:\n        return None  # prompt alone exceeds context: reject, don't truncate silently\n    # An unset cap means \"up to the context limit\", not \"small\".\n    requested = requested_max_output if requested_max_output is not None else remaining\n    return min(requested, remaining)\n```\n\n**Takeaway:** Treat an unspecified `max_tokens`\n\nas the context maximum, because that's exactly what it can become in production.\n\nAdmission 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.\n\n``` python\nimport threading\n\nclass KVAdmissionController:\n    def __init__(self, total_kv_bytes, per_token_bytes, safety_fraction=0.9):\n        self.usable = int(total_kv_bytes * safety_fraction)  # leave headroom\n        self.per_token = per_token_bytes\n        self.reserved = 0\n        self.lock = threading.Lock()\n\n    def try_admit(self, prompt_tokens, output_cap):\n        need = (prompt_tokens + output_cap) * self.per_token\n        with self.lock:\n            if self.reserved + need <= self.usable:\n                self.reserved += need\n                return need  # a handle to release later\n            return None      # rejected: no budget\n\n    def release(self, need):\n        with self.lock:\n            self.reserved = max(0, self.reserved - need)\n```\n\nTwo details that bite people. First, `release`\n\nmust run in a `finally`\n\nblock — 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`\n\nis 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.\n\n**Takeaway:** Hold the reservation at full size for the whole request and release it in `finally`\n\n— a leaked reservation is a slow-motion outage.\n\n`Retry-After`\n\nmatter so much?\nA 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`\n\nwith a `Retry-After`\n\nheader telling the client how long to wait.\n\n``` python\nfrom fastapi import FastAPI, Request\nfrom fastapi.responses import JSONResponse\n\napp = FastAPI()\n\n@app.post(\"/generate\")\nasync def generate(req: Request):\n    body = await req.json()\n    prompt_tokens = count_tokens(body[\"prompt\"])\n    cap = effective_output_cap(body.get(\"max_tokens\"), prompt_tokens)\n    if cap is None:\n        return JSONResponse({\"error\": \"prompt exceeds context window\"}, status_code=413)\n\n    handle = controller.try_admit(prompt_tokens, cap)\n    if handle is None:\n        return JSONResponse(\n            {\"error\": \"server at capacity\", \"retry_after_seconds\": 2},\n            status_code=429,\n            headers={\"Retry-After\": \"2\"},\n        )\n    try:\n        return await run_generation(body, prompt_tokens, cap)\n    finally:\n        controller.release(handle)\n```\n\n`Retry-After`\n\nworks 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`\n\nbeats 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.\"\n\n**Takeaway:** A rejection without `Retry-After`\n\ndoesn't shed load, it reschedules it for one second from now — the header is what actually protects you.\n\n| Strategy | What it counts | Fails when | Backpressure |\n|---|---|---|---|\n| Concurrency cap (max N requests) | Request count only | Long-context requests, mixed sizes | Usually none |\n| Prompt-length admission | Prompt tokens | Short prompt, large `max_tokens`\n|\nOften none |\n| Prompt + output reservation | Prompt + generation cap | Mostly holds; over-reserves early stops |\n`429` + `Retry-After`\n|\n\nA 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.\n\n**Should I reserve KV cache for max_tokens or the actual output length?**\n\nReserve for `max_tokens`\n\n(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.\n\n**What should an LLM server return when it's out of KV cache?**\n\nReturn HTTP `429 Too Many Requests`\n\nwith a `Retry-After`\n\nheader. 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.\n\n**Why does a short prompt cause more trouble than a long one?**\n\nBecause 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`\n\nreserves for 4,200 tokens; a 3,000-token prompt with `max_tokens=200`\n\nreserves for 3,200. The short one is the bigger promise.\n\nIf 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`\n\nblock. When there's no budget, reject with `429`\n\nand a `Retry-After`\n\nheader 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.", "url": "https://wpnews.pro/news/reserve-kv-cache-for-the-output-you-promised-not-just-the-prompt-you-received", "canonical_source": "https://dev.to/libme/reserve-kv-cache-for-the-output-you-promised-not-just-the-prompt-you-received-oko", "published_at": "2026-08-12 03:08:53+00:00", "updated_at": "2026-08-12 03:14:53.418792+00:00", "lang": "en", "topics": ["large-language-models", "ai-infrastructure", "mlops", "ai-safety"], "entities": [], "alternates": {"html": "https://wpnews.pro/news/reserve-kv-cache-for-the-output-you-promised-not-just-the-prompt-you-received", "markdown": "https://wpnews.pro/news/reserve-kv-cache-for-the-output-you-promised-not-just-the-prompt-you-received.md", "text": "https://wpnews.pro/news/reserve-kv-cache-for-the-output-you-promised-not-just-the-prompt-you-received.txt", "jsonld": "https://wpnews.pro/news/reserve-kv-cache-for-the-output-you-promised-not-just-the-prompt-you-received.jsonld"}}