Skip to content

Memory budgeting: auto-size to the workload by default; optional percent/absolute cap with a pre-flight error instead of an OOM #83

Description

@localai-bot

One of the sharpest usability complaints about vLLM is that you have to compute your own VRAM budget, express it as a percentage, and get it right — or the engine OOMs. The number depends on the model, the quant, the context length, the concurrency and whatever else is already resident on the card, and none of that is anything a user should be doing arithmetic about.

We should not mirror that ergonomics. This issue proposes the shape we want instead:

  • No option given → the engine works out what the declared run needs and allocates exactly that. Nothing to tune.
  • Option given (percent or an absolute size) → that is a hard cap on the engine's total device footprint, enforced by a pre-flight check that produces a precise, actionable error rather than an OOM.

Where we are

We are currently behind vLLM on this axis, not ahead of it. There is no memory profiling at all — the KV pool is a raw block count the user types in:

  • EngineParams::num_blocks = 256 — "KV blocks to allocate" (include/vllm/entrypoints/model_loader.h:58), alongside block_size = 32 (:57), max_model_len (:59) and max_num_seqs = 8 (:60).
  • The server exposes it verbatim as --num-blocks N (examples/server/main.cpp:100,203-204,370).
  • The C ABI carries it as vllm_model_params.num_blocks, defaulted to the same 256 (src/capi/vllm_c.cpp:429,485).
  • It lands as BlockPool(num_gpu_blocks, ...) (include/vllm/v1/core/block_pool.h:96,223), which asserts > 0 and otherwise trusts it (src/vllm/v1/core/block_pool.cpp:51).

So today a user has to convert "I have 40 GB free and want 32k context at concurrency 8" into a block count by hand. That is strictly worse than a percentage.

The seam is already rowed as owed work: .agents/feature-matrix.md:83 carries Sizing: gpu_memory_utilization, block overrides as PARTIAL T0 — "watermark/fixed loader inputs exist; public utilization/cache-byte/override policy absent" — pointing at a specs/kv-sizing.md that was never written.

What upstream already has

.agents/porting-inventory.md:89 records the three upstream knobs, all in config/cache.py, all T0:

knob semantics
gpu_memory_utilization fraction of total device memory the engine may use; default 0.9
kv_cache_memory_bytes absolute KV-pool size, bypassing the fraction
num_gpu_blocks_override pin the pool to an exact block count

vLLM profiles a forward pass and derives available KV memory as total × utilization − non-torch − peak activation, erroring when that lands at or below zero.

That machinery is worth mirroring, but it does not solve the complaint, for three reasons:

  1. The fraction is of total, not free. On any card that is not exclusively yours, the fraction you need is a function of what someone else is already holding, which is exactly the arithmetic the user is being asked to do.
  2. Weights load before the knob engages. The utilization figure sizes the KV pool after the model is resident, so an oversized model OOMs during load and never reaches the check.
  3. 0.9 is taken whether or not it is needed. A 4B model on an 80 GB card reserves 72 GB it will never touch, for no benefit, and blocks anything else on the device.

Proposed design

Mode 1 — default: size to the declared workload

With no memory flag, the engine computes required bytes per allocation class, before allocating anything:

class source
weights checkpoint metadata (safetensors header / GGUF manifest) — known before reading a byte of tensor data
CUDA context measured once at context creation
peak activation profile run at max_num_batched_tokens
KV cache max_model_len × max_num_seqs, at the resolved block_size and KV dtype
CUDA-graph pools capture-set footprint

It then allocates exactly that and leaves the remainder of the device free. This is the surpass over vLLM: upstream takes its 90% regardless of whether the workload needs 8 GiB or 80.

If the total exceeds free memory, it fails at startup — see Mode 3.

Mode 2 — a cap on the total engine footprint

A limit caps everything the engine allocates on device — weights, activations, KV, graph pools, context — not just the KV pool. This distinction is the whole point: a cap that only governs KV cannot prevent the weight-load OOM, which is the failure users actually hit.

Three spellings, all capping the same total:

  • --memory-limit 40GiB — absolute size string. Unambiguous, and the only form that is safe on a unified pool (see below).
  • --gpu-memory-utilization 0.85 — vLLM's exact flag name and fraction semantics, so existing vLLM launch lines port unchanged (per the mirror-vLLM policy).
  • --num-gpu-blocks-override N — upstream's reproducibility escape hatch; pins the KV pool to an exact block count regardless of profiling. This is where today's --num-blocks goes: demoted from primary knob to explicit override.

Precedence between the three must be spelled out explicitly and tested, not left to argument order.

Mode 3 — refuse before allocating, with a usable error

When the requirement exceeds the cap (or free memory in Mode 1), the engine must exit before allocating, printing the accounting and concrete remedies:

error: engine requires 83.9 GiB, limit is 40.0 GiB (short by 43.9 GiB)

  weights          61.2 GiB   (bf16, 39 layers)
  activations       3.1 GiB   (peak @ max_num_batched_tokens=4096)
  KV cache         18.4 GiB   (max_model_len=32768 x max_num_seqs=8, bf16)
  cuda graphs       1.2 GiB
  cuda context      2.0 GiB

  to fit within 40.0 GiB, try one of:
    --max-model-len 8192      (KV -> 4.6 GiB)
    --max-num-seqs 2          (KV -> 4.6 GiB)
    --kv-cache-dtype fp8      (KV -> 9.2 GiB)
    a smaller quantization    (weights dominate this budget)

The remedies must be computed from the actual budget, not printed as a generic hint. "You are 43.9 GiB over and here are the three levers that close it" is the difference between a usable error and a stack trace.

Unified-memory hazard (GB10 and friends)

On GB10 the ~119 GiB pool is unified: a fraction-of-total setting reserves host RAM as well, and gpu_memory_utilization=0.85 has hard-rebooted our DGX three separate times. This is not a hypothetical.

Two consequences for the design:

  1. Absolute bytes is the primary form, with the percentage kept for vLLM compatibility rather than as the recommended interface.
  2. The accounting must know whether the pool is unified, which makes this a Platform-seam question — free/total memory and an is-unified predicate belong behind the platform abstraction (ROAD-V1-C1), not in a CUDA-specific branch. Discrete-GPU and unified-pool devices need different safety margins, and hard-coding either is wrong.

Milestones

  • M0, spike — write .agents/specs/kv-sizing.md (already planned: at .agents/feature-matrix.md:83, never written). Inventory config/cache.py and the upstream memory-profiling path with file:line anchors; define the per-class accounting model and the precedence rules.
  • M1, accounting — a MemoryBudget that computes required bytes per class without allocating, plus the Platform seam for free/total device memory and the is-pool-unified predicate. Unit-gated against known checkpoints (predicted weight bytes == actual resident bytes).
  • M2, auto-sizing default — derive the block count from the budget; --num-blocks becomes --num-gpu-blocks-override. Gate: every existing model gate runs with no block flag and stays token-exact.
  • M3, the caps--memory-limit, --gpu-memory-utilization, --num-gpu-blocks-override, with documented precedence, wired through the server flags and the C ABI (vllm_model_params). Gate: at a matched --gpu-memory-utilization, our KV pool matches vLLM's own for the same model and config.
  • M4, the pre-flight refusal — refuse-before-allocate with the breakdown and computed remedies. Gate: a deliberately over-subscribed config exits cleanly on GB10 (non-zero exit, no OOM, no box reboot), covering the unified-pool path specifically.
  • M5, runtime guard (optional) — an allocation that would exceed the budget mid-run fails the request, not the engine.

Docs (README, docs/STATUS.md) update in the same change as the milestone that shifts externally-visible behaviour, per the keep-README-current rule.

Correctness

Memory sizing changes the KV pool size, which changes preemption and scheduling timing but not emitted tokens. The SACRED token-exact gates should be unaffected, and M2's gate makes that explicit by re-running them with no block flag at all.


Roadmap row: ROAD-V1-MEM in .agents/roadmap_v1.md.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions