VRAM Admission Control: Reject Agent Requests Before OOM

The OOM killer is a post-mortem, not a scheduler. By the time the kernel or the CUDA allocator kills a process, the request that caused it has already consumed the GPU. Admission control moves that decision forward: decide whether a request can fit before it touches VRAM.

This matters most for agent swarms. A single chat completion has a predictable footprint. An agent loop does not. It spawns tool calls, retries, re-reads context, and holds KV cache for every parallel branch. Teams typically observe that the failure mode is not one giant request but ten medium ones arriving within the same second.

Why does the OOM killer fail agent workloads?

The OOM killer is reactive and blind to intent. It fires when memory is already exhausted, and it picks a victim by heuristic, not by priority. In a swarm, the victim is often a long-running generation whose KV cache represents minutes of work.

The result is a cascade: one eviction frees memory, the scheduler admits queued work, and the next wave hits the same wall. Retries amplify load exactly when the system is least able to absorb it. A rejected request, by contrast, is cheap. The client learns immediately and can back off.

What does a pre-flight admission check need to know?

Admission control needs an estimate of the request's peak footprint, not its current one. Peak VRAM for a transformer inference request is roughly weights (shared, counted once) plus KV cache for the full context plus activation scratch. KV cache scales with layers × heads × head_dim × tokens × bytes_per_element × concurrency.

A workable estimator takes the request's declared max tokens and concurrency and returns a reservation:

def estimate_kv_bytes(cfg, tokens, batch):
    per_token = (
        2                      # key and value
        * cfg.n_layers
        * cfg.n_kv_heads
        * cfg.head_dim
        * cfg.dtype_bytes
    )
    return per_token * tokens * batch

Two rules keep this honest. First, reserve against max_tokens, not the prompt length, because generation grows the cache. Second, add a fixed scratch margin for activations and fragmentation; allocators rarely hand back every freed block.

How do you track reservations without a full allocator?

You do not need to mirror the CUDA allocator. You need a reservation ledger: a single source of truth for committed VRAM, decremented on completion and timeout. Every admit is a compare-and-swap against the ledger.

def try_admit(ledger, req, capacity, margin):
    need = estimate_kv_bytes(req.cfg, req.max_tokens, req.batch) + margin
    with ledger.lock:
        if ledger.committed + need > capacity:
            return Reject(reason="insufficient_vram", need=need)
        ledger.committed += need
        return Admit(reservation_id=ledger.new_id(need))

The hard part is leak prevention. A crashed worker that never releases its reservation poisons the ledger. Tie every reservation to a lease with a deadline, and have a reaper reclaim expired leases. Treat the ledger as authoritative and the actual allocator as a check, not the reverse.

Should you reject, queue, or degrade?

Rejection is one option among three, and the right choice depends on the caller.

• Reject for interactive requests. A fast 429 with a retry hint beats a request that stalls for thirty seconds and then dies.
• Queue for batch and background agents. Bounded queues with a max wait, so backlog cannot grow without limit.
• Degrade for anything that tolerates it: shorter max_tokens, smaller batch, or a quantized path. Degradation must be explicit in the response, not silent.

The decision belongs in policy, not scattered across handlers. Encode it as a table keyed by request class, and let operators change the table without a redeploy.

Before you add another retry layer, instrument the ledger: log committed versus capacity and the rejection reason for every refused request. That single time series tells you whether you are capacity-bound or estimate-bound.

How does admission control interact with swarm coordination?

A swarm has a coordination layer that fans out work. If that layer is unaware of VRAM pressure, it will happily spawn twenty branches against a GPU that fits four. The fix is to expose admission state to the coordinator as a backpressure signal, not just a per-request error.

When the ledger crosses a high-water mark, the coordinator should stop fanning out and serialize branches. When it drops, resume. This turns a hard cliff into a slope. The coordinator does not need to know the exact byte count; it needs a boolean or a coarse level (green, amber, red) that it can act on.

One subtlety: fan-out decisions and admissions race. Two branches can both read green and both get rejected. That is fine if rejection is cheap. It is not fine if the coordinator treats rejection as a fatal error. Design the contract so a rejected branch is a normal, retryable outcome.

What are the failure modes to design against?

Admission control adds a new component that can itself fail. Plan for these:

• Stale estimates. Model config changes shift KV size. Version the estimator with the model config.
• Ledger drift. If the ledger and the allocator disagree, trust the ledger for admission and alert on the gap.
• Thundering herd. A mass rejection triggers synchronized retries. Add jitter and a retry-after hint.
• Priority inversion. A low-priority batch job holding a reservation blocks an interactive request. Preempt by lease, not by hope.

The common thread: every mechanism needs an escape hatch and an observable. An admission layer you cannot debug is worse than the OOM killer it replaced.

If you run agent swarms on shared GPUs, start by logging peak VRAM per request class for a week, then set your first reservation margin from that distribution. What does your current scheduler do when a request cannot fit: reject fast, queue blind, or let the allocator decide?