[Edge] Size the KV cache from the model, not from a server default

`VLLM_CPU_KVCACHE_SPACE` is a server-shaped default and is parsed with
`int()`, so it cannot express an edge-sized cache at all -- the smallest
non-zero value it accepts is 1 GiB. Measured on Spark-X2.5-1.7B (Xeon 8480C,
2048-token single stream), passing `kv_cache_memory_bytes` sized from the
model instead of the 4 GiB default takes resident memory from 8329 MB to
5254 MB, a 3.1 GB saving with no effect on output.

`kv_budget_for` reads the config and prices the cache. It reports two numbers
because the model and the backend disagree: a hybrid-attention model does not
need a full-length cache in every layer -- Spark is 21 sliding (window 512)
to 7 full layers, so at 32 k context its real cache is 3.8x smaller than
`layers x context` -- but vLLM's CPU path does not take that discount today
(for a 4 GiB budget it reported 71 527 tokens, i.e. the flat 56 KiB/token).
So `bytes_total` is the flat figure, which is what is safe to pass now, and
`bytes_total_hybrid` is reported beside it for a backend that accounts per
layer type.

Deliberately conservative in three places, because a budget that is too small
is not a graceful degradation -- vLLM preempts by recompute, which on a
single-stream edge deployment is a stall: the window is rounded up to the
block size, a 1.25 headroom factor is applied, and a config that has a
`sliding_window` but no `layer_types` map is charged the flat rate rather than
assumed to apply the window everywhere.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

tests/edge/test_kv_budget.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Model-sized KV budgets, including hybrid attention."""
+
+import pytest
+
+from vllm_omni.edge.kv_budget import dtype_bytes, kv_budget_for
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Cfg:
+    def __init__(self, **kw):
+        self.__dict__.update(kw)
+
+
+def _spark(**kw) -> _Cfg:
+    """Spark-X2.5-1.7B: 28 layers, 3 sliding (512) to 1 full, 2 KV heads, 256."""
+    layer_types = ["sliding_attention"] * 3 + ["full_attention"]
+    base = dict(
+        num_hidden_layers=28, num_attention_heads=8, num_key_value_heads=2,
+        head_dim=256, hidden_size=2048, sliding_window=512,
+        layer_types=layer_types * 7,
+    )
+    base.update(kw)
+    return _Cfg(**base)
+
+
+def test_flat_bytes_per_token_matches_the_hand_calculation():
+    # 28 layers x 2 KV heads x 256 head_dim x 2 (K and V) x 2 bytes = 56 KiB.
+    b = kv_budget_for(_spark(), max_model_len=2048)
+    assert b.bytes_per_token_flat == 28 * 2 * 256 * 2 * 2
+    assert b.bytes_per_token_flat == 56 * 1024
+
+
+def test_hybrid_attention_is_charged_for_its_real_cache():
+    b = kv_budget_for(_spark(), max_model_len=32768)
+    assert (b.full_layers, b.sliding_layers, b.window) == (7, 21, 512)
+    assert b.token_layers_hybrid == 7 * 32768 + 21 * 512
+    assert b.token_layers_flat == 28 * 32768
+    # The recommended budget stays flat, because vLLM's CPU accounting is flat.
+    assert b.bytes_total > b.bytes_total_hybrid
+    assert b.saving_vs_flat == pytest.approx(917504 / 240128, rel=1e-3)
+
+
+def test_a_short_context_never_pays_for_a_longer_window():
+    """A 512-window layer at a 128-token context caches 128, not 512."""
+    b = kv_budget_for(_spark(), max_model_len=128)
+    assert b.token_layers_hybrid == 28 * 128
+
+
+def test_budget_scales_with_sequences_and_dtype():
+    one = kv_budget_for(_spark(), max_model_len=2048, max_num_seqs=1)
+    four = kv_budget_for(_spark(), max_model_len=2048, max_num_seqs=4)
+    assert four.bytes_total == 4 * one.bytes_total
+    fp8 = kv_budget_for(_spark(), max_model_len=2048, kv_dtype="fp8")
+    assert fp8.bytes_total * 2 == one.bytes_total
+
+
+def test_edge_sized_budget_is_far_under_the_server_default():
+    """The measured point: 2048 tokens needs ~112 MiB, not the 4 GiB default."""
+    b = kv_budget_for(_spark(), max_model_len=2048)
+    mib = b.bytes_total / 2**20
+    assert 130 < mib < 150, mib          # flat 56 KiB/token x 2048 x 1.25
+    assert b.bytes_total < 4 * 2**30 / 25  # vs the 4 GiB default
+    assert b.bytes_total_hybrid / 2**20 < 70
+
+
+def test_plain_model_without_layer_types_uses_a_flat_budget():
+    cfg = _Cfg(num_hidden_layers=32, num_attention_heads=32,
+               num_key_value_heads=8, hidden_size=4096)
+    b = kv_budget_for(cfg, max_model_len=4096)
+    assert b.sliding_layers == 0 and b.window is None
+    assert b.token_layers_hybrid == b.token_layers_flat
+    assert b.saving_vs_flat == 1.0
+    assert b.bytes_total == b.bytes_total_hybrid
+    # head_dim inferred from hidden_size / heads.
+    assert b.bytes_per_token_flat == 32 * 8 * 128 * 2 * 2
+
+
+def test_sliding_window_without_layer_types_is_not_assumed():
+    """A config with a window but no per-layer map might apply it everywhere or
+    nowhere; charging the smaller cache would risk preemption."""
+    cfg = _Cfg(num_hidden_layers=4, num_attention_heads=8, num_key_value_heads=8,
+               head_dim=64, sliding_window=256)
+    b = kv_budget_for(cfg, max_model_len=4096)
+    assert b.sliding_layers == 0
+    assert b.token_layers_hybrid == b.token_layers_flat
+
+
+def test_unreadable_config_raises_rather_than_guessing():
+    with pytest.raises(ValueError, match="num_hidden_layers"):
+        kv_budget_for(_Cfg(), max_model_len=2048)
+    with pytest.raises(ValueError, match="head_dim"):
+        kv_budget_for(_Cfg(num_hidden_layers=4), max_model_len=2048)
+
+
+def test_dtype_bytes_covers_the_kv_dtypes_vllm_accepts():
+    assert dtype_bytes("bfloat16") == 2
+    assert dtype_bytes("float32") == 4
+    assert dtype_bytes("fp8") == 1
+    assert dtype_bytes("something_unknown") == 2  # conservative default
+
+
+def test_summary_names_the_shape_it_priced():
+    s = kv_budget_for(_spark(), max_model_len=32768).summary()
+    assert "flat" in s and "21 sliding@512" in s and "MiB" in s

vllm_omni/edge/kv_budget.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Size the KV cache from the model instead of from a server default.
+
+``VLLM_CPU_KVCACHE_SPACE`` defaults to a server-shaped budget and is parsed
+with ``int()``, so it cannot even express an edge-sized cache -- the smallest
+non-zero value it accepts is 1 GiB. Measured on Spark-X2.5-1.7B (Xeon 8480C,
+2048-token single stream), asking for what the model actually needs instead of
+the 4 GiB default takes resident memory from **8329 MB to 5254 MB**.
+
+A hybrid-attention model does not *need* a full-length cache in every layer.
+Spark-X2.5 is 3 sliding layers (window 512) to 1 full layer over 28 layers, so
+at a 32 k context its real cache is 3.8x smaller than a flat budget:
+
+    flat      28 x 32768               = 917 504 token-layers
+    hybrid    21 x 512 + 7 x 32768     = 240 128 token-layers
+
+But vLLM's CPU path does not currently take that discount: for a 4 GiB budget
+it reported 71 527 tokens, i.e. 56 KiB/token, which is the flat figure. So
+``bytes_total`` is the **flat** number -- the one that is safe to pass today --
+and the hybrid number is reported alongside it as ``bytes_total_hybrid``, for
+a backend that accounts per layer type. A budget that is too small is not a
+graceful degradation: vLLM preempts by recompute, which on a single-stream
+edge deployment is a stall.
+"""
+
+from dataclasses import dataclass
+from typing import Any
+
+_DTYPE_BYTES = {
+    "float32": 4, "fp32": 4,
+    "bfloat16": 2, "bf16": 2, "float16": 2, "fp16": 2, "half": 2,
+    "float8_e4m3fn": 1, "fp8": 1, "int8": 1,
+}
+
+
+def dtype_bytes(dtype: Any) -> int:
+    name = getattr(dtype, "name", None) or str(dtype).replace("torch.", "")
+    return _DTYPE_BYTES.get(name, 2)
+
+
+@dataclass(frozen=True)
+class KVBudget:
+    bytes_total: int
+    """Flat budget: safe to pass to vLLM today."""
+    bytes_total_hybrid: int
+    """What the model's own cache needs, if the backend accounts per layer type."""
+    bytes_per_token_flat: int
+    token_layers_hybrid: int
+    token_layers_flat: int
+    sliding_layers: int
+    full_layers: int
+    window: int | None
+
+    @property
+    def saving_vs_flat(self) -> float:
+        """What a layer-type-aware backend would save; 1.0 when there is none."""
+        return (
+            self.token_layers_flat / self.token_layers_hybrid
+            if self.token_layers_hybrid
+            else 1.0
+        )
+
+    def summary(self) -> str:
+        extra = ""
+        if self.sliding_layers:
+            extra = (
+                f"; {self.bytes_total_hybrid / 2**20:.0f} MiB if the backend "
+                f"accounts {self.sliding_layers} sliding@{self.window} layers "
+                f"separately ({self.saving_vs_flat:.2f}x)"
+            )
+        return f"{self.bytes_total / 2**20:.0f} MiB flat{extra}"
+
+
+def kv_budget_for(
+    hf_config: Any,
+    *,
+    max_model_len: int,
+    max_num_seqs: int = 1,
+    kv_dtype: Any = "bfloat16",
+    block_size: int = 16,
+    headroom: float = 1.25,
+) -> KVBudget:
+    """Bytes of KV cache a deployment of this model actually needs.
+
+    Reads ``layer_types``/``sliding_window`` when present, so hybrid-attention
+    models are charged for their real cache and not for the longest layer.
+    """
+    n_layers = int(
+        getattr(hf_config, "num_hidden_layers", None)
+        or getattr(hf_config, "n_layer", 0)
+    )
+    if n_layers <= 0:
+        raise ValueError("cannot read num_hidden_layers from the config")
+    n_heads = int(getattr(hf_config, "num_attention_heads", 0) or 0)
+    kv_heads = int(getattr(hf_config, "num_key_value_heads", 0) or n_heads or 1)
+    head_dim = int(
+        getattr(hf_config, "head_dim", 0)
+        or (getattr(hf_config, "hidden_size", 0) // max(n_heads, 1))
+    )
+    if head_dim <= 0:
+        raise ValueError("cannot read head_dim from the config")
+
+    layer_types = list(getattr(hf_config, "layer_types", None) or [])
+    window = getattr(hf_config, "sliding_window", None)
+    sliding = sum(1 for t in layer_types if "sliding" in str(t)) if window else 0
+    full = n_layers - sliding
+
+    def rounded(tokens: int) -> int:
+        return ((tokens + block_size - 1) // block_size) * block_size
+
+    per_seq_full = rounded(max_model_len)
+    # A sliding layer still has to hold the window, and never more than the
+    # context: a 512-window layer at a 128-token context caches 128.
+    per_seq_sliding = rounded(min(int(window), max_model_len)) if sliding else 0
+
+    token_layers_hybrid = max_num_seqs * (
+        full * per_seq_full + sliding * per_seq_sliding
+    )
+    token_layers_flat = max_num_seqs * n_layers * per_seq_full
+    bytes_per_token_layer = kv_heads * head_dim * 2 * dtype_bytes(kv_dtype)
+
+    return KVBudget(
+        bytes_total=int(token_layers_flat * bytes_per_token_layer * headroom),
+        bytes_total_hybrid=int(token_layers_hybrid * bytes_per_token_layer * headroom),
+        bytes_per_token_flat=n_layers * bytes_per_token_layer,
+        token_layers_hybrid=token_layers_hybrid,
+        token_layers_flat=token_layers_flat,
+        sliding_layers=sliding,
+        full_layers=full,
+        window=int(window) if sliding else None,
+    )