This document is relevant for: Trn2, Trn3

determine_available_memory design#

1) KV cache budget#

By default, vLLM Neuron uses the memory-utilization and fractional-cap policy:

logical_core_hbm = hbm_stack_bytes / logical_cores_per_hbm_stack
total_budget = int(logical_core_hbm * gpu_memory_utilization)
user_budget = max(total_budget - model_weight_bytes, 0)
heuristic_cap = int(total_budget * cap_fraction)
available = max(min(user_budget, heuristic_cap), 0)

The default cap_fraction is 0.30. The worker respects gpu_memory_utilization and VLLM_NEURON_KV_GMU_BUDGET_CAP_FRACTION. This conservative policy limits the KV cache to a fraction of the memory-utilization budget.

An explicit headroom value enables the reserve policy:

available = max(logical_core_hbm - model_weight_bytes - headroom_bytes, 0)

hbm_stack_bytes is the HBM capacity reported for one physical stack.
logical_cores_per_hbm_stack is 1 at LNC 2 and 2 at LNC 1.
model_weight_bytes includes model parameters and buffers.
headroom_bytes reserves HBM per logical core for all other allocations.

The headroom includes internal tensors, scratchpad, DMA rings, runtime data, compiled code, and a safety margin.

Only headroom mode sets gpu_memory_utilization to 1.0 and bypasses the fractional cap. An explicit zero also enables headroom mode.

2) Why headroom is necessary#

The compiler includes the KV cache shape in each compiled program. A larger KV cache changes the graph and its memory requirements.

The compiler gives the best estimate of non-KV memory after it receives a KV shape. However, the worker must select the KV shape before compilation. The worker therefore uses a configurable estimate for non-weight and non-KV memory.

In CPU inference, the default policy uses:

total_budget = int(host_available / local_workers * gpu_memory_utilization)
available = int(total_budget * cap_fraction)

Explicit headroom uses:

available = max(host_available / local_workers - headroom_bytes, 0)

Host available memory already excludes loaded weights and current allocations, so CPU mode does not subtract model weights again.

The reserve has no default. For example, 4 GiB is a tested reserve for GPT-OSS 120B on Trn3 with the Native backend and MXFP4 weights. Validate the reserve for your model and configuration.

3) Configuration#

Variable	Default	Description
`VLLM_NEURON_KV_GMU_BUDGET_CAP_FRACTION`	`0.30`	Fractional cap for default sizing. Headroom mode bypasses this value.
`VLLM_NEURON_RESERVED_HEADROOM_GIB`	unset	Explicit reserve per worker. Enables headroom sizing on both routes.
`VLLM_NEURON_MIN_KV_BUDGET_GIB`	`1.0`	Minimum KV budget. The server stops if the budget is smaller.

In headroom mode, increase the active reserve if model loading or compilation reports NCC_EOOM002, Maximum peak HBM usage, NRT_RESOURCE in nrt_load_util, or Allocation Failure. First confirm that the same configuration passes with a smaller KV cache or with the default policy. Start with 1 GiB increments. Each additional 1 GiB of headroom reduces the KV cache budget by 1 GiB. Rerun the same model, sharding, backend, LNC degree, and hardware configuration, and keep the smallest value that passes. Each backend deployment can use a different value for the same headroom variable.

Do not increase headroom when the worker reports that the KV cache budget is below its minimum threshold. Decrease the active headroom variable for that error. A validated configuration can also decrease headroom to use a larger KV cache. At LNC 1, calculate the limit from half of the HBM stack:

maximum_headroom =
    hbm_stack_bytes / 2 - model_weight_bytes - minimum_kv_budget

The startup log reports the logical-core HBM, model weights, active headroom, and resulting KV budget. An override at or above the logical-core HBM share leaves no usable KV cache.

Both backends and CPU use VLLM_NEURON_RESERVED_HEADROOM_GIB. When this variable is unset, they use default fractional-cap sizing. CPU compilation and device execution select the same budget policy.

In default mode, reduce the fractional cap to reduce KV capacity after an HBM allocation failure. An explicit headroom value changes the budget policy; validate the resulting KV capacity before deployment.

4) Implementation flow#

The implementation is in vllm_neuron/vllm/worker/neuron_worker.py.

determine_available_memory()
  |
  +-- VLLM_NEURON_CPU_COMPILE == 1
  |      total_hbm = get_total_available_memory()
  |      model_weight_bytes = _get_byte_used_from_model()
  |      return _compute_kv_budget(total_hbm, model_weight_bytes)
  |
  +-- VLLM_NEURON_CPU_MODE == 1
  |      use the selected policy with the host memory fair share
  |
  +-- Neuron device
         total_hbm = runtime_bytes_used + runtime_bytes_free
         model_weight_bytes = _get_byte_used_from_model()
         return _compute_kv_budget(total_hbm, model_weight_bytes)

CPU compilation and device execution use the same model weight calculation. This keeps the KV block count equal and permits compilation cache hits.

_compute_kv_budget(hbm_stack_bytes, model_weight_bytes):
    logical_core_hbm = (
        hbm_stack_bytes // get_logical_cores_per_hbm_stack()
    )
    if headroom is explicit:
        available = max(
            logical_core_hbm - model_weight_bytes - headroom_bytes, 0
        )
    else:
        total_budget = int(logical_core_hbm * gpu_memory_utilization)
        available = max(min(
            max(total_budget - model_weight_bytes, 0),
            int(total_budget * cap_fraction),
        ), 0)
    if available < VLLM_NEURON_MIN_KV_BUDGET_GIB: raise RuntimeError
    return available

The device path uses runtime statistics only to get total HBM. It does not use runtime memory usage as the model weight size.

The CPU inference path uses current available host memory, which already reflects model weights and other allocations. Neither CPU inference policy subtracts model weights again.

HBM stack sharing (LNC 1) scales the total first. One HBM stack serves 2 physical NeuronCores. At LNC 2 a logical core owns the stack. At LNC 1 two logical cores share it. Both inputs still report the whole stack. The runtime reports the stack to each core that holds a part of it, and the static table is per stack. get_logical_cores_per_hbm_stack() gives the divisor, which is 1 at LNC 2 and 2 at LNC 1. Without it, two siblings each size a KV cache against the same memory, and the allocation fails at load (nrt_tensor_allocate status=4). The scaling is in _compute_kv_budget, so both paths that call it get the same budget. This keeps the KV block count equal, which the compile cache needs.
CPU Compile only mode uses the static HBM size for the target platform (looked up via NEURON_PLATFORM_TARGET_OVERRIDE) and computes bytes_used from model parameter and buffer sizes. No device is present — the goal is to produce the same KV block count that the execution phase will use, ensuring compilation cache hits.
Neuron mode (compile/execute) also uses _get_byte_used_from_model() (model parameter + buffer sizes after model loading) rather than raw runtime.get_vnc_memory_stats() for bytes_used. This ensures the KV block count is identical between CPU compile and on-device execution. The runtime stats are only used to derive total_hbm_bytes.
CPU mode (no compilation) is a separate path that uses host memory fair-share — it does not go through _compute_kv_budget and is unrelated to the compile/execute flow.
An explicit headroom value reserves non-KV memory after the LNC share is known.
Minimum-KV guard fails fast for unusable configs.

5) Minimum budget check#

The worker compares both CPU and Neuron results with VLLM_NEURON_MIN_KV_BUDGET_GIB. It stops early if the result is too small.

In default mode, the error names memory utilization and the fractional cap. In headroom mode, it names the active reserve and reports its maximum value. Increasing headroom makes a minimum-KV-budget error worse.

Known HBM allocation exceptions raised while loading, initializing the KV cache, compiling, or warming up retain their original exception and add a note with the active reserve and the 1 GiB tuning procedure in headroom mode.