[Edge][CPU] An explicit KV budget must beat the legacy env var
`CpuPlatform.check_and_update_config` ends with a block whose own comment
reads "Lagecy setting":
env_key = "VLLM_CPU_KVCACHE_SPACE"
if env_key in os.environ and os.environ[env_key] != "":
cache_config.kv_cache_memory_bytes = int(os.environ[env_key]) * GiB
It runs unconditionally, so `--kv-cache-memory-bytes` is silently discarded
whenever the env var is set -- which it has to be on CPU, or the engine tries
to claim most of the machine. Two consequences for edge serving: the budget
can only be expressed in whole GiB (parsed with `int()`, so the smallest
non-zero cache is 1 GiB), and vllm-omni's per-stage `kv_cache_memory_bytes`,
which is how the deploy YAMLs express per-stage budgets, has no effect at all
on CPU.
Verified by allocation rather than by memory, because allocation cannot drift:
with the env at 1 GiB, asking for 128 MiB produced the same 17 881 cache
tokens as asking for nothing, and with the env at 4 GiB it produced 71 527 --
the explicit value never applied. After restoring the obvious precedence
(explicit wins, env stays the default for callers that set nothing) the same
request yields 2 231 tokens, the 128 MiB it asked for.
The Qwen3-TTS CPU overlay now sizes stage 0 from the model -- 112 KiB/token,
so 2 seqs at 2048 tokens needs 560 MiB, not 4 GiB. Stage 1 is left alone on
purpose: Code2Wav is a generation stage whose cache is not a plain attention
KV, and under-sizing makes vLLM preempt by recompute, which on a single-stream
edge deployment is a stall.
What this saves in resident memory is deliberately *not* claimed. VmRSS on
this host moves 1.2-2.0 GB between identical runs, because it counts the
mmap'd checkpoint's file pages; an earlier 3.1 GB figure came from a
comparison that changed the model file as well as the budget and does not
stand. RssAnon is the metric to use and that measurement is still in flight.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
5f27f60bde
commit
ca4dd105e4
2 files changed
+76
-1
No files matched your search
@@ -155,7 +155,21 @@ platforms:
|
||||
max_model_len: 2048
|
||||
gpu_memory_utilization: 0.05
|
||||
dtype: bfloat16
|
||||
kv_cache_memory_bytes: 4294967296
|
||||
# Sized from the model rather than from a server default: the
|
||||
# talker is 28 layers x 8 KV heads x 128 head_dim x 2 (K,V) x 2 B
|
||||
# = 112 KiB/token, so max_num_seqs 2 at max_model_len 2048 needs
|
||||
# 560 MiB against the 4 GiB this used to ask for -- 7.3x the need.
|
||||
# (vllm_omni/edge/kv_budget.py computes it from the HF config.)
|
||||
#
|
||||
# This only takes effect with the patch in vllm_omni/patch.py:
|
||||
# CpuPlatform.check_and_update_config overwrites
|
||||
# kv_cache_memory_bytes from VLLM_CPU_KVCACHE_SPACE unconditionally,
|
||||
# so on CPU this key was dead. Verified by allocation: the same
|
||||
# request went from 71 527 cache tokens to 2 231 once the explicit
|
||||
# value won. What that saves in resident memory is *not* established
|
||||
# -- RssAnon is the metric to use, since VmRSS here moves 1-2 GB
|
||||
# between identical runs as the mmap'd checkpoint's pages come and go.
|
||||
kv_cache_memory_bytes: 671088640
|
||||
runtime:
|
||||
devices: "cpu"
|
||||
env:
|
||||
@@ -169,6 +183,10 @@ platforms:
|
||||
max_model_len: 8192
|
||||
gpu_memory_utilization: 0.05
|
||||
dtype: bfloat16
|
||||
# Left at 2 GiB deliberately: Code2Wav is a generation stage and its
|
||||
# cache geometry is not a plain attention KV, so the figure cannot be
|
||||
# derived the way stage 0's can. Measure before shrinking it --
|
||||
# under-sizing makes vLLM preempt by recompute, which is a stall.
|
||||
kv_cache_memory_bytes: 2147483648
|
||||
runtime:
|
||||
devices: "cpu"
|
||||
|
||||
@@ -637,3 +637,60 @@ def _patch_cpu_aot_compile_cache_load() -> None:
|
||||
|
||||
|
||||
_patch_cpu_aot_compile_cache_load()
|
||||
|
||||
|
||||
def _patch_cpu_explicit_kv_cache_memory() -> None:
|
||||
"""Stop a legacy env var from overriding an explicit KV budget on CPU.
|
||||
|
||||
``CpuPlatform.check_and_update_config`` ends with (vLLM 0.28, the comment
|
||||
is theirs)::
|
||||
|
||||
# Lagecy setting
|
||||
env_key = "VLLM_CPU_KVCACHE_SPACE"
|
||||
if env_key in os.environ and os.environ[env_key] != "":
|
||||
cache_config.kv_cache_memory_bytes = int(os.environ[env_key]) * GiB
|
||||
|
||||
-- unconditionally, so ``--kv-cache-memory-bytes`` is silently discarded
|
||||
whenever the env var is set, which it must be on CPU or the engine tries to
|
||||
claim most of the machine. Verified: with the env at 1 GiB, asking for
|
||||
128 MiB produced the same 17 881 tokens as asking for nothing, and with the
|
||||
env at 4 GiB it produced 71 527 -- the explicit value never applied.
|
||||
|
||||
Two consequences for edge serving, which is where KV size actually matters:
|
||||
the budget can only be set in whole GiB (the env is parsed with ``int()``,
|
||||
so the smallest non-zero cache is 1 GiB), and a deploy config that sets
|
||||
``kv_cache_memory_bytes`` per stage -- which is how vllm-omni expresses
|
||||
per-stage budgets -- has no effect at all on CPU.
|
||||
|
||||
This restores the obvious precedence: an explicit budget wins, and the env
|
||||
remains the default for callers that set nothing. Spark-X2.5 needs 56 KiB
|
||||
per token, so a 2048-token single-stream deployment wants 112 MiB where the
|
||||
env's floor is 1 GiB.
|
||||
"""
|
||||
from vllm.platforms.cpu import CpuPlatform
|
||||
|
||||
if getattr(CpuPlatform, "_omni_explicit_kv_bytes", False):
|
||||
return
|
||||
|
||||
base = CpuPlatform.check_and_update_config.__func__
|
||||
|
||||
def check_and_update_config(cls, vllm_config):
|
||||
cache_config = getattr(vllm_config, "cache_config", None)
|
||||
explicit = getattr(cache_config, "kv_cache_memory_bytes", None)
|
||||
base(cls, vllm_config)
|
||||
if explicit is None or cache_config is None:
|
||||
return
|
||||
if cache_config.kv_cache_memory_bytes != explicit:
|
||||
cache_config.kv_cache_memory_bytes = explicit
|
||||
_PATCH_LOGGER.info(
|
||||
"Restored explicit kv_cache_memory_bytes=%d (%.3f GiB); "
|
||||
"VLLM_CPU_KVCACHE_SPACE would have overridden it.",
|
||||
explicit,
|
||||
explicit / (1024**3),
|
||||
)
|
||||
|
||||
CpuPlatform.check_and_update_config = classmethod(check_and_update_config)
|
||||
CpuPlatform._omni_explicit_kv_bytes = True
|
||||
|
||||
|
||||
_patch_cpu_explicit_kv_cache_memory()
|
||||
Reference in new issue
Block a user