rebase and move kv cache group validation

Signed-off-by: qizixi <[email protected]>

Author: zixi-qi

vllm/v1/spec_decode/eagle.py

@@ -14,6 +14,7 @@
 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
+from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.sample.metadata import SamplingMetadata
 
 logger = init_logger(__name__)
@@ -352,6 +353,18 @@ def dummy_run(
                 hidden_states=self.hidden_states[:num_tokens],
             )
 
+    def validate_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
+        kv_cache_groups: dict[str, int] = {}
+        for id, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
+            for layer_name in kv_cache_group.layer_names:
+                kv_cache_groups[layer_name] = id
+        assert len(
+            set([
+                kv_cache_groups[layer_name]
+                for layer_name in self.attn_layer_names
+            ])) == 1, "For multi-layer eagle draft model, "
+        "all layers should belong to the same kv cache group"
+
 
 # NOTE(woosuk): Currently, the below code is not used and we always use argmax
 # to sample the draft tokens. We will use this after we find a way to manage

vllm/v1/worker/gpu_model_runner.py

@@ -1890,12 +1890,10 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
         self.initialize_attn_backend(kv_cache_config)
 
         kv_caches: dict[str, torch.Tensor] = {}
-        kv_cache_group_ids: dict[str, int] = {}
 
-        for id, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
+        for i, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
             kv_cache_spec = kv_cache_group.kv_cache_spec
             for layer_name in kv_cache_group.layer_names:
-                kv_cache_group_ids[layer_name] = id
                 tensor_config = kv_cache_config.tensors[layer_name]
                 assert tensor_config.size % kv_cache_spec.page_size_bytes == 0
                 num_blocks = tensor_config.size // kv_cache_spec.page_size_bytes
@@ -1922,12 +1920,9 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
 
         if self.speculative_config and self.speculative_config.use_eagle():
             assert isinstance(self.drafter, EagleProposer)
-            assert len(
-                set([
-                    kv_cache_group_ids[layer_name]
-                    for layer_name in self.drafter.attn_layer_names
-                ])) == 1, "For multi-layer eagle draft model, "
-            "all layers should belong to the same kv cache group"
+            # validate all draft model layers belong to the same kv cache
+            # group
+            self.drafter.validate_kv_cache_group(kv_cache_config)
 
         bind_kv_cache(
             kv_caches,