From ca4dd105e4f4c334db6eb102a04ca833cf809299 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Sat, 12 Sep 2026 14:16:01 +0800 Subject: [PATCH] [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) --- vllm_omni/deploy/qwen3_tts.yaml | 20 +++++++++++- vllm_omni/patch.py | 57 +++++++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index d7cc489e2..f8d3a17e0 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -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" diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index e2cd479ce..91ebd1b90 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -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()