An inference node that runs out of GPU memory does not fail gracefully. It panics. The CUDA runtime raises a cudaErrorMemoryAllocation error, the process typically crashes or enters an unrecoverable state, and every in-flight request on that node is dropped. If you are running multiple nodes behind a load balancer, some traffic continues to be served while the affected node restarts, but the requests that were mid-generation are lost and must be retried by the caller.
For short-context workloads, this is an edge case that appears under unusual load patterns and resolves quickly on restart. For long-context workloads serving 32K, 64K, or 128K token contexts, it is a routine operational problem. The KV cache grows linearly with context length, and a node that was operating comfortably at 8K contexts can run out of memory with no notice when a sudden batch of 64K requests arrives. The consequences of an OOM under these conditions are proportionally worse: each dropped request represents tens of seconds of prefill compute that must be re-run from scratch.
This post covers how we handle pre-OOM detection, request migration, and post-pressure recovery in Inferact's scheduler. These mechanisms were built directly in response to production incidents where our test deployments lost significant compute to avoidable OOM panics.
Memory Pressure Signals Before OOM
The standard GPU memory allocation APIs provide current allocated bytes and total capacity, but they do not provide a direct warning that you are approaching the limit. NVIDIA's Management Library (NVML) exposes memory statistics through nvmlDeviceGetMemoryInfo, which returns used, free, and total bytes. We poll this at 500 ms intervals per node.
The challenge is that this polling approach has two weaknesses. First, 500 ms is a long time during a prefill phase for a 32K-context request. The KV cache for such a request grows from zero to approximately 8 GB (for Llama-3 8B in FP16 at 32K tokens) over the course of the prefill. If we poll every 500 ms, we can miss the transition from safe to critical memory state within a single polling interval. Second, NVML's GetMemoryInfo reports allocated bytes, not reserved bytes. CUDA maintains a memory pool that retains freed allocations for reuse. The reserved pool can be substantially larger than the currently allocated tensors, which means the "free" bytes reported by NVML overstate actual available allocation headroom.
We address the first weakness by supplementing NVML polling with a per-request KV budget tracker. When a request is admitted, the scheduler estimates its peak KV cache consumption based on the input token count and an estimated output token budget. This estimate is subtracted from the node's available KV headroom before the request starts execution. If admitting a new request would bring the projected KV headroom below a safety threshold (currently 15 percent of total KV pool capacity), the request is queued rather than admitted, even if the NVML query shows available memory.
We address the second weakness by periodically forcing CUDA pool garbage collection via torch.cuda.empty_cache() on each node and re-measuring the post-GC free memory. We run this every 30 seconds and use the post-GC value as our baseline for safety threshold calculations.
The Three Pressure Levels
We model memory pressure as three levels, each triggering different scheduler behaviors.
Level 1 (caution): Projected KV headroom is between 15 percent and 25 percent of total KV pool capacity. At this level, the scheduler stops admitting new requests to this node but continues allowing in-flight requests to complete. New requests that would have gone to this node are redirected to other nodes or held in queue if no other nodes have capacity.
Level 2 (pressure): Projected KV headroom drops below 15 percent. At this level, the scheduler begins preemptive KV cache eviction of the least-recently-accessed blocks that are not part of currently active decode sequences. This is aggressive LRU eviction intended to free KV pool space before the physical memory limit is reached. In-flight requests continue; queue admission remains paused.
Level 3 (critical): Post-GC free memory is below 2 GB on an 80 GB node (approximately 2.5 percent of capacity). At this level, the scheduler marks the node as unavailable for new admissions and begins migrating queued requests. In-flight requests that have not yet entered prefill are transferred to alternate nodes; requests already in prefill are allowed to complete but new slots are not opened on this node.
The 2 GB threshold for Level 3 is not derived from a universal principle. It reflects the minimum free memory we have observed to be safe on our A100 nodes before an OOM panic occurs during a large-batch prefill. Nodes with smaller VRAM (such as RTX 4090 at 24 GB) use a proportionally smaller absolute threshold but a higher percentage threshold (approximately 8 percent of total VRAM).
Request Migration Mechanics
Migrating a queued request (one that has not yet started prefill) is straightforward: the scheduler removes it from the current node's admission queue and re-inserts it into the admission queue for an alternate node. The request's full context is in the caller's request payload, so no inter-node state transfer is required.
Migrating a request that has already begun prefill is much harder. The partial KV cache state built up during prefill is stored in the current node's GPU memory. There is no native mechanism to serialize and transfer this state to another node over the network during live prefill execution. In practice, we do not migrate in-progress prefill requests. Instead, when Level 3 is reached, the scheduler marks the prefilling requests as requiring restart: they are placed back into the global queue with their original input context intact, and the partial prefill work is discarded. The caller does not see an error; they see a delayed response as the request restarts on a new node. The restart cost is the prefill recompute for whatever fraction of the context had been processed before the migration.
We are not satisfied with this approach. Discarding partial prefill work is expensive, especially for 64K-context requests where prefill can take several seconds. The right solution is a mechanism to checkpoint and restore partial prefill state across nodes, which requires either a high-bandwidth network transfer (feasible on IB-connected clusters) or a shared memory tier accessible from both nodes (NVMe-backed swap or CPU DRAM). Neither is currently implemented in our scheduler. The current approach is a pragmatic fallback that avoids OOM panics at the cost of occasional expensive restarts.
Post-Pressure Recovery
After a node reaches Level 3 and migrates queued requests, the scheduler does not immediately restore the node to full admission. We run a recovery sequence:
Step 1: Force CUDA pool GC and measure post-GC free memory. If the node is still below the Level 2 threshold after GC, we wait for in-flight sequences to complete and their KV pages to be released.
Step 2: Once post-GC free memory returns above the Level 2 threshold, the node transitions to Level 1 status. At this point, the scheduler restores single-request admission: the node accepts one new request at a time, waiting for post-admission memory measurement before accepting the next.
Step 3: After 5 consecutive successful single-request admissions without triggering pressure indicators, the node returns to full admission status and normal batch scheduling resumes.
The step-by-step re-admission is conservative. It adds latency before the node is fully utilized again, which reduces throughput during the recovery period. We have considered more aggressive recovery strategies but opted for conservatism after experiencing two incidents where premature full re-admission triggered a second OOM panic on a node that had not fully drained its in-flight long-context requests.
What Long Context Changes About This Problem
Most of the memory management techniques described in inference literature were developed for workloads where context lengths are modest (2K-8K tokens). At these lengths, the KV cache per request is small, the memory pressure signals give you meaningful advance warning, and migration or preemption is cheap because partial prefill represents a small amount of discarded work.
At 64K or 128K token contexts, every assumption changes. Prefill for a 128K-token request on Llama-3 8B takes 10-30 seconds on a single A100, depending on batch occupancy. A single such request consumes approximately 32 GB of KV cache in FP16 (128K tokens, 32 layers, GQA with 8 KV heads, 128 head dimension, 2 bytes per element). That is 40 percent of an A100's VRAM, consumed by a single request. The memory pressure model built for 8K contexts does not transfer cleanly to this regime.
We have not fully solved long-context memory management. The partial-prefill migration problem is real and the current restart approach is a workaround, not a solution. Teams running production workloads at 64K+ contexts should be aware that the graceful degradation story for memory pressure at those lengths is less mature than it is for shorter contexts. We are actively working on this, and the next substantive change will likely be adding CPU DRAM offloading for KV state as a buffer during migration, which avoids the full prefill restart cost at the price of slower decode for migrated requests.