diff --git a/tests/edge/test_kv_budget.py b/tests/edge/test_kv_budget.py new file mode 100644 index 000000000..85c1af22c --- /dev/null +++ b/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 diff --git a/vllm_omni/edge/kv_budget.py b/vllm_omni/edge/kv_budget.py new file mode 100644 index 000000000..195a68d75 --- /dev/null +++ b/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, + )