Admission Control for Self-Hosted LLMs: Rejecting Requests Before the KV Cache OOMs You

#ai #llm #backend #performance

The short version: an LLM server that loads fine and answers your test prompt will still crash under real traffic, because GPU memory for inference is dominated by the KV cache — per-request memory that grows with context length and concurrency, not by the model weights. The fix that keeps a deployment alive isn't a bigger card; it's admission control: deciding, before you accept a request, whether there is enough KV-cache budget to finish it. This post is how to compute that budget and where to enforce it.

I'm writing this as a follow-up after a commenter on an earlier post about GPU memory made a sharp point: the first deploy should be budgeted around admission control, not just whether the weights load, and a single long-context request is often a better stress test than a pile of short chats. That's exactly right, and it deserves its own walkthrough.

Why does a server that "fits" still crash under load?

Loading the weights proves the model fits. It says nothing about whether the server serves. During inference, the memory that actually decides survival is the KV cache: the cached key/value tensors for every token already in a request's context, so the model doesn't recompute them each decode step. Its size per request is roughly:

kv_bytes_per_request =
    2 × num_layers × num_kv_heads × head_dim × seq_len × bytes_per_element

The 2 is keys plus values. Note num_kv_heads, not the total attention-head count — modern models use grouped-query attention (GQA), which shrinks this by 4–8× versus the naive formula. The term that hurts is seq_len, and it's multiplied again by however many requests you run concurrently. Weights are a fixed one-time cost; the KV cache is a per-request cost that you pay again for every simultaneous user and every extra token of context.

That's why the failure mode is nonlinear. Eight short chats might sit comfortably. One 6,000-token document at a batch of eight can blow past the same budget instantly, because each of those eight requests now carries a large KV footprint at once.

Takeaway: Weights tell you the model loads; the KV cache tells you how many concurrent tokens you can actually hold — and that's the number that decides whether you stay up.

How do I compute my KV-cache budget?

Start from what's left after the weights and runtime overhead, then divide by the per-token cost. Here's a calculator you can run before you rent anything:

def kv_bytes_per_token(num_layers, num_kv_heads, head_dim, dtype_bytes=2):
    # 2 = keys + values. dtype_bytes: 2 for FP16/BF16 KV, 1 for FP8/INT8 KV.
    return 2 * num_layers * num_kv_heads * head_dim * dtype_bytes

def token_budget(total_gib, weight_gib, overhead_gib, per_token_bytes):
    free_bytes = (total_gib - weight_gib - overhead_gib) * (1024 ** 3)
    if free_bytes <= 0:
        raise ValueError("No memory left for KV cache after weights + overhead")
    return int(free_bytes // per_token_bytes)

# Llama-3-8B-class: 32 layers, 8 KV heads (GQA), head_dim 128, BF16 KV
per_tok = kv_bytes_per_token(num_layers=32, num_kv_heads=8, head_dim=128)
budget = token_budget(total_gib=24, weight_gib=16, overhead_gib=2, per_token_bytes=per_tok)

print(f"{per_tok/1024:.1f} KiB per token")      # ~128 KiB/token
print(f"~{budget:,} total KV tokens in flight") # the number that matters

The output — total KV tokens you can hold at once — is your real capacity. If it says ~48,000 tokens, that's six requests at 8k context, or forty-eight at 1k, or one giant 48k-context request that leaves room for nothing else. Concurrency and context length trade against the same pool.

Takeaway: Your capacity isn't "N requests"; it's a fixed number of KV tokens in flight, and every request spends from it in proportion to its context length.

Where should I reject: the API edge or the scheduler?

Both, at different jobs. The commenter's question — cap context at the edge, or let the scheduler reject once the cache budget is gone — has a "yes, and" answer: the edge does cheap, predictable rejection; the scheduler does the real-time backpressure the edge can't see.

Layer	What it checks	Rejects with	Good at	Blind to
API edge (your code)	`input_tokens + max_new_tokens` vs a hard cap	HTTP 400 (too long) / 429 (busy)	Cheap, instant, predictable	Live GPU occupancy right now
Concurrency gate	In-flight request/token count vs budget	HTTP 429 + `Retry-After`	Coarse load shedding	Exact per-request KV cost
Inference scheduler (vLLM/TGI)	Actual free KV blocks per step	Queue, preempt, or recompute	True memory state	Client's latency SLO

A hard context cap at the edge is the single highest-value control, because an unbounded-context request is the one that turns a healthy server into an OOM. Reject it with a clear 400 before it ever reaches the GPU. Then a concurrency gate sheds coarse load with a 429 so the scheduler is never asked to do the impossible. Finally, the scheduler (vLLM's PagedAttention allocator, or TGI's max-concurrent-requests limit) handles the fine-grained, moment-to-moment truth. As of mid-2026, vLLM will even preempt and recompute a running request when it runs short on KV blocks — correct, but it costs latency, which is why you want the upstream gates absorbing most of the pressure.

Takeaway: Cap context at the edge for correctness, gate concurrency for load, and let the scheduler own the last-millisecond truth — don't ask any one layer to do all three.

What does a minimal admission gate look like?

Here's an edge gate that enforces a context cap and a token-in-flight budget before proxying to your inference backend. It's deliberately simple — a semaphore over a token count — because the goal is to fail fast and cheap:

import asyncio
from fastapi import FastAPI, HTTPException, Request

app = FastAPI()

MAX_CONTEXT = 8192          # hard per-request cap: input + generation
KV_TOKEN_BUDGET = 48_000    # from token_budget() above, with headroom
_in_flight = 0
_lock = asyncio.Lock()

async def admit(cost_tokens: int):
    global _in_flight
    async with _lock:
        if _in_flight + cost_tokens > KV_TOKEN_BUDGET:
            raise HTTPException(429, "KV budget exhausted",
                                headers={"Retry-After": "2"})
        _in_flight += cost_tokens
    return cost_tokens

async def release(cost_tokens: int):
    global _in_flight
    async with _lock:
        _in_flight -= cost_tokens

@app.post("/generate")
async def generate(req: Request):
    body = await req.json()
    prompt_tokens = body["prompt_tokens"]          # count upstream, not len(text)
    max_new = min(body.get("max_new_tokens", 512), MAX_CONTEXT - prompt_tokens)
    cost = prompt_tokens + max_new
    if cost > MAX_CONTEXT:
        raise HTTPException(400, f"context {cost} exceeds cap {MAX_CONTEXT}")

    await admit(cost)
    try:
        return await call_backend(body)   # your vLLM/TGI call
    finally:
        await release(cost)

Two things trip people up here. First, count prompt_tokens with the model's real tokenizer upstream — len(text) is not a token count and will let a request in that the GPU can't actually hold. Second, reserve for max_new_tokens, not just the prompt: generation grows the KV cache one token per step, so a short prompt with a huge generation budget is still an expensive request. Budget for the worst case you admitted, not the prompt you received.

Takeaway: Admit on prompt_tokens + max_new_tokens measured by the real tokenizer, and always return a Retry-After on 429 so clients back off instead of hammering.

FAQ

How do I calculate KV cache size for an LLM?
Per token, it's 2 × num_layers × num_kv_heads × head_dim × dtype_bytes. Multiply by sequence length for one request, and sum across all concurrent requests to get total KV memory in use. Use num_kv_heads (not total heads) for GQA models, or you'll overestimate by several times.

Should I return 429 or 400 when an LLM request is too big?
Use 400 when a single request exceeds your hard context cap — it will never succeed, so the client must change it. Use 429 with a Retry-After header when the request is valid but the server is momentarily out of KV budget — the same request will succeed once load drops.

Why does my LLM server OOM under load but not in testing?
Testing usually sends one short prompt, which costs almost no KV cache. Real traffic sends long contexts at concurrency, and KV-cache memory scales with context_length × concurrent_requests. The weights fit the whole time; the KV cache is what overflows.

Bottom line

Size your deployment by KV-token budget, not by whether the weights load. Put a hard context cap at the API edge so a single unbounded request can't take down the box, add a concurrency gate that sheds load with a 429 and Retry-After, and let your inference server's scheduler own the last-millisecond allocation. If you're on a serving stack today, vLLM is the one that already does fine-grained KV paging and preemption for you — but it serves best when an upstream admission gate keeps it out of the impossible cases. Do the arithmetic before you rent the GPU, and your first deploy survives contact with real users.