A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 02:15 UTC
[v1][core] Support for attention free models (#20811)
Signed-off-by: Christian Pinto <christian.pinto@ibm.com>
Christian Pinto committed
Jul 15, 2025 at 15:20 UTC
4ffd963fa07c35c7945c2b5ae30577d00e18f80d
3 files changed
+33
-3
vllm/v1/core/kv_cache_manager.py
+6
-1
@@ -78,7 +78,12 @@ class KVCacheManager:
78
) -> None:
79
self.max_model_len = max_model_len
80
81
+ if len(kv_cache_config.kv_cache_groups) == 0:
82
+ # Attention free models don't have kv cache,
83
+ # thus don't need prefix caching.
84
+ enable_caching = False
85
self.enable_caching = enable_caching
86
+
87
self.caching_hash_fn = (
88
sha256_cbor_64bit if caching_hash_algo == "sha256_cbor_64bit" else
89
sha256 if caching_hash_algo == "sha256" else hash)
@@ -101,7 +106,7 @@ class KVCacheManager:
106
kv_cache_config=kv_cache_config,
107
max_model_len=self.max_model_len,
108
use_eagle=self.use_eagle,
104
- enable_caching=enable_caching,
109
+ enable_caching=self.enable_caching,
110
caching_hash_fn=self.caching_hash_fn,
111
enable_kv_cache_events=enable_kv_cache_events,
112
)
vllm/v1/core/kv_cache_utils.py
+20
-1
@@ -563,6 +563,10 @@ def check_enough_kv_cache_memory(vllm_config: VllmConfig,
563
ValueError: If there is not enough memory available for the KV cache.
564
"""
565
566
+ # No need to check for available memory if the kv_cache_spec is empty
567
+ if not kv_cache_spec:
568
+ return
569
+
570
if available_memory <= 0:
571
raise ValueError("No available memory for the cache blocks. "
572
"Try increasing `gpu_memory_utilization` when "
@@ -749,6 +753,13 @@ def is_kv_cache_page_size_uniform(
753
return len(page_sizes) == 1
754
755
756
+def is_kv_cache_type_attention_free(
757
+ kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
758
+
759
+ # kv_cache_spec is an empty dict for attention free models
760
+ return not kv_cache_spec
761
+
762
+
763
def _get_kv_cache_config_uniform_page_size(
764
vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec],
765
available_memory: int) -> KVCacheConfig:
@@ -891,6 +902,10 @@ def _get_kv_cache_config_uniform_page_size(
902
return kv_cache_config
903
904
905
+def _get_kv_cache_config_attention_free() -> KVCacheConfig:
906
+ return KVCacheConfig(num_blocks=1, kv_cache_tensors=[], kv_cache_groups=[])
907
+
908
+
909
def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
910
"""
911
This function tries to convert the KV cache specs to one type if the model
@@ -957,7 +972,11 @@ def get_kv_cache_config(
972
if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
973
unify_hybrid_kv_cache_specs(kv_cache_spec)
974
960
- if is_kv_cache_type_uniform(kv_cache_spec):
975
+ if is_kv_cache_type_attention_free(kv_cache_spec):
976
+ # This returns a kv_cache config with 0 kv_cache groups and 1 block
977
+ # to allow for the KVCache manager to handle attention free models.
978
+ return _get_kv_cache_config_attention_free()
979
+ elif is_kv_cache_type_uniform(kv_cache_spec):
980
# KV cache of all layers are the same, which is true for
981
# most models. Allocate the same amount of memory for
982
# each layer.
vllm/v1/engine/core.py
+7
-1
@@ -139,7 +139,13 @@ class EngineCore:
139
140
# Profiles the peak memory usage of the model to determine how much
141
# memory can be allocated for kv cache.
142
- available_gpu_memory = self.model_executor.determine_available_memory()
142
+ has_kv_cache = any(kv_cache_spec for kv_cache_spec in kv_cache_specs)
143
+ if has_kv_cache:
144
+ available_gpu_memory = \
145
+ self.model_executor.determine_available_memory()
146
+ else:
147
+ # Attention free models don't need memory for kv cache
148
+ available_gpu_memory = [0] * len(kv_cache_specs)
149
150
assert len(kv_cache_specs) == len(available_gpu_memory)
151
# Get the kv cache tensor size