Support for attention free models revisited to reuse existing KVCache maanger.

Signed-off-by: Christian Pinto <[email protected]>

Author: christian-pinto

vllm/v1/core/kv_cache_coordinator.py

@@ -219,7 +219,9 @@ def __init__(self, kv_cache_config: KVCacheConfig, max_model_len: int,
         super().__init__(kv_cache_config, max_model_len, use_eagle,
                          enable_caching, caching_hash_fn,
                          enable_kv_cache_events)
-        self.verify_and_split_kv_cache_groups()
+        # attention free models are initialized with 0 kv_cache_groups
+        if len(self.kv_cache_config.kv_cache_groups) > 0:
+            self.verify_and_split_kv_cache_groups()
 
     def verify_and_split_kv_cache_groups(self) -> None:
         """

vllm/v1/core/kv_cache_manager.py

@@ -84,12 +84,17 @@ def __init__(
         self.log_stats = log_stats
         # FIXME: make prefix cache stats conditional on log_stats
         self.prefix_cache_stats = PrefixCacheStats() if log_stats else None
-        assert len(
-            set(g.kv_cache_spec.block_size
-                for g in kv_cache_config.kv_cache_groups)
-        ) == 1, "Only one block size is supported for now"
-        self.block_size = kv_cache_config.kv_cache_groups[
-            0].kv_cache_spec.block_size
+
+        if len(kv_cache_config.kv_cache_groups) == 0:
+            #This is an attention free model that is started with 0 KVCache groups.
+            self.block_size = 0
+        else:
+            assert len(
+                set(g.kv_cache_spec.block_size
+                    for g in kv_cache_config.kv_cache_groups)
+            ) == 1, "Only one block size is supported for now"
+            self.block_size = kv_cache_config.kv_cache_groups[
+                0].kv_cache_spec.block_size
 
         self.coordinator = get_kv_cache_coordinator(
             kv_cache_config=kv_cache_config,

vllm/v1/core/kv_cache_utils.py

@@ -551,6 +551,10 @@ def check_enough_kv_cache_memory(vllm_config: VllmConfig,
         ValueError: If there is not enough memory available for the KV cache.
     """
 
+    # No need to check for available memory if the model is attention free
+    if vllm_config.model_config.is_attention_free:
+        return
+
     if available_memory <= 0:
         raise ValueError("No available memory for the cache blocks. "
                          "Try increasing `gpu_memory_utilization` when "
@@ -736,6 +740,11 @@ def is_kv_cache_page_size_uniform(
     page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()}
     return len(page_sizes) == 1
 
+def is_kv_cache_type_attention_free(
+        kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
+
+    if "attention_free" in kv_cache_spec:
+        return True
 
 def _get_kv_cache_config_uniform_page_size(
         vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec],
@@ -877,6 +886,11 @@ def _get_kv_cache_config_uniform_page_size(
     return kv_cache_config
 
 
+def _get_kv_cache_config_attention_free() -> KVCacheConfig:
+    return KVCacheConfig(num_blocks=1,
+                         kv_cache_tensors=[],
+                         kv_cache_groups=[])
+
 def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
     """
     This function tries to convert the KV cache specs to one type if the model
@@ -943,7 +957,9 @@ def get_kv_cache_config(
     if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
         unify_hybrid_kv_cache_specs(kv_cache_spec)
 
-    if is_kv_cache_type_uniform(kv_cache_spec):
+    if is_kv_cache_type_attention_free(kv_cache_spec):
+        return _get_kv_cache_config_attention_free()
+    elif is_kv_cache_type_uniform(kv_cache_spec):
         # KV cache of all layers are the same, which is true for
         # most models. Allocate the same amount of memory for
         # each layer.

vllm/v1/engine/core.py

@@ -134,33 +134,22 @@ def _initialize_kv_caches(
             self, vllm_config: VllmConfig) -> tuple[int, int, KVCacheConfig]:
         start = time.time()
 
-        #TODO: CP start from here
-        if vllm_config.model_config.is_attention_free:
-            # No need for initializing anything related to KV cache if the model
-            # is attention free.
-            kv_cache_specs = []
-            kv_cache_configs = [
-                KVCacheConfig(num_blocks=0,
-                              kv_cache_tensors=[],
-                              kv_cache_groups=[])
-            ]
-        else:
-            # Get all kv cache needed by the model
-            kv_cache_specs = self.model_executor.get_kv_cache_specs()
-
-            # Profiles the peak memory usage of the model to determine how much
-            # memory can be allocated for kv cache.
-            available_gpu_memory = (
-                self.model_executor.determine_available_memory())
-
-            assert len(kv_cache_specs) == len(available_gpu_memory)
-            # Get the kv cache tensor size
-            kv_cache_configs = [
-                get_kv_cache_config(vllm_config, kv_cache_spec_one_worker,
-                                    available_gpu_memory_one_worker)
-                for kv_cache_spec_one_worker, available_gpu_memory_one_worker
-                in zip(kv_cache_specs, available_gpu_memory)
-            ]
+        # Get all kv cache needed by the model
+        kv_cache_specs = self.model_executor.get_kv_cache_specs()
+
+        # Profiles the peak memory usage of the model to determine how much
+        # memory can be allocated for kv cache.
+        available_gpu_memory = (
+            self.model_executor.determine_available_memory())
+
+        assert len(kv_cache_specs) == len(available_gpu_memory)
+        # Get the kv cache tensor size
+        kv_cache_configs = [
+            get_kv_cache_config(vllm_config, kv_cache_spec_one_worker,
+                                available_gpu_memory_one_worker)
+            for kv_cache_spec_one_worker, available_gpu_memory_one_worker
+            in zip(kv_cache_specs, available_gpu_memory)
+        ]
 
         # Since we use a shared centralized controller, we need the
         # `kv_cache_config` to be consistent across all workers to make sure

vllm/v1/executor/abstract.py

@@ -73,10 +73,16 @@ def register_failure_callback(self, callback: FailureCallback):
         pass
 
     def determine_available_memory(self) -> list[int]:  # in bytes
+        if self.vllm_config.model_config.is_attention_free:
+            return [0]
+
         output = self.collective_rpc("determine_available_memory")
         return output
 
     def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
+        if self.vllm_config.model_config.is_attention_free:
+            return [{"attention_free": KVCacheSpec(block_size=0)}]
+
         output = self.collective_rpc("get_kv_cache_spec")
         return output