[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