SIGN IN SIGN UP

[Bugfix][KV Offload] Bound sparse KV offload memory by NPU, DRAM and workload limits (#14921)

### What this PR does / why we need it?

Sparse KV offload (SFA with `enable_sparse_li_c8`) had two problems in
memory planning:

1. `SparseKVOffloadManager` always sized the CPU KV pool at the **full
configured `dram_size_per_dp_GB`**, regardless of how many KV blocks the
engine actually planned to use. On large-DRAM configs (e.g. 180 GiB per
DP group) this allocates far more DRAM than needed and makes pool
initialization very slow (tens of minutes of hugepage-backed
allocation), even when the active workload only needs a fraction of it.
2. The worker-side logic
(`update_available_memory_for_sparse_kv_offload`) only enlarged total
available memory by a host/device page-size ratio. The number of KV
blocks was therefore never bounded by DRAM capacity or by active
workload demand, and when the blocks implied by free NPU memory exceeded
the configured DRAM, startup failed with a generic "please increase
`sparse_kv_offload_config.dram_size_per_dp_GB`" error.

This PR replaces that logic with an explicit memory plan:

- New `plan_sparse_kv_offload_memory()` returns a
`SparseKVOffloadMemoryBudget` that bounds the block count by the minimum
of three limits and reports the limiting factor:
- **NPU capacity** — free device memory (counting the full page size
when `keep_device_kv_cache=True`, device portion only otherwise);
- **DRAM limit** — `dram_size_per_dp_GB` minus a per-layer alignment
reserve (each host layer may lose up to 3 × 2 MiB to 2 MiB-aligned
allocations);
- **Workload demand** — `max_blocks_per_request × max_num_seqs` plus the
null block.
- `NPUWorker.determine_available_memory()` now applies the budget
through `_apply_kv_offload_decode_memory_constraints()` and returns the
planned bytes, with a one-shot log line showing every limit, the final
block count, the limiting factor, and the host/device/alignment byte
breakdown.
- `SparseKVOffloadManager` sizes the actual CPU pool from the planned
block count (`num_blocks × host page sizes + alignment reserve`, via the
new `get_sparse_kv_offload_cpu_pool_size_bytes()`) instead of the full
configured DRAM, and raises an actionable error (planned vs. limit vs.
`num_blocks`) if the planned pool still exceeds the DRAM limit.

As a result, the offload pool is allocated only as large as needed: DRAM
footprint and CPU pool initialization time drop accordingly, and
undersized `dram_size_per_dp_GB` configurations fail early with a clear
message.

### Does this PR introduce _any_ user-facing change?

Yes. With sparse KV offload enabled:

- The CPU offload pool is no longer always allocated at the full
`dram_size_per_dp_GB`; it is sized from the planned block count, so the
number of KV cache blocks can be smaller than before when DRAM or
workload demand is the limiting factor. The limiting factor and full
plan are logged at startup.
- The previous error "Needed dram size ... Please increase
`sparse_kv_offload_config.dram_size_per_dp_GB`" is replaced by a new
"planned CPU pool exceeds DRAM limit after alignment" error carrying
planned/limit/num_blocks values.
- Log messages around sparse KV offload memory planning have changed.

### How was this patch tested?

Unit tests added:

-
`tests/ut/distributed/kv_transfer/sparse_kv_offload/test_sparse_kv_offload_manager.py`
— memory plan limited by workload / DRAM / NPU (each asserting block
counts, planned bytes and the limiting factor); non-positive NPU or DRAM
capacity producing zero blocks; invalid spec layouts rejected (missing
host or device specs, mixed block sizes); CPU pool size including the
per-layer alignment reserve; `UniformTypeKVCacheSpecs` support; manager
initializing the backend with the planned pool size (rank 0 allocates,
other ranks 0) with a TP barrier; and rejection when the planned pool
exceeds the DRAM limit.
- `tests/ut/worker/a2/test_worker_v1.py` — constraints disabled returns
memory unchanged; planner budget returned by
`_apply_kv_offload_decode_memory_constraints`; `kv_cache_spec` fetched
when not cached; explicit `kv_cache_memory_bytes` path and the normal
`determine_available_memory` path both apply the constraints.

- vLLM version: v0.27.1
- vLLM main:
https://github.com/vllm-project/vllm/commit/ba07e4a48fc951300d97eb506217dd530583dea3

---------

Signed-off-by: Abstrey <846134527@qq.com>
Z
Zangzizhe committed
6b52b112e63109e124abd75dc2f67edb769874b6
Parent: d72770f
Committed by GitHub <noreply@github.com> on 8/28/2026, 2:27:08 AM