[V1] Support cross-layer KV sharing

Signed-off-by: Yong Hoon Shin <[email protected]>

Author: sarckk

tests/v1/core/test_kv_cache_utils.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+import copy
 import importlib
 
 import pytest
@@ -15,12 +16,13 @@
                                          PrefixCachingMetrics,
                                          estimate_max_model_len,
                                          generate_block_hash_extra_keys,
+                                         get_kv_cache_config,
                                          hash_block_tokens,
                                          hash_request_tokens,
                                          unify_kv_cache_configs)
 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
-                                        KVCacheGroupSpec, KVCacheTensor,
-                                        SlidingWindowSpec)
+                                        KVCacheGroupSpec, KVCacheSpec,
+                                        KVCacheTensor, SlidingWindowSpec)
 from vllm.v1.metrics.stats import PrefixCacheStats
 from vllm.v1.request import Request
 
@@ -557,13 +559,76 @@ def test_merge_kv_cache_spec():
     assert merged_layer_spec.sliding_window == 1
 
 
+def test_get_kv_cache_config_cross_layer_kv_sharing():
+    # Create a VllmConfig
+    model_id = "Qwen/Qwen1.5-7B"
+    max_model_len = 16383
+    model_config = ModelConfig(
+        model_id,
+        task="generate",
+        tokenizer=model_id,
+        tokenizer_mode="auto",
+        trust_remote_code=False,
+        seed=0,
+        dtype="float16",
+        max_model_len=max_model_len,
+    )
+    scheduler_config = SchedulerConfig(max_num_batched_tokens=32768)
+
+    vllm_config = VllmConfig(
+        model_config=model_config,
+        scheduler_config=scheduler_config,
+    )
+
+    # Create KV cache specs
+
+    # max memory usage bytes calculated as:
+    # 1024 * 2 * 16 * 32 * 128 * 2
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=32,
+        head_size=128,
+        dtype=torch.float16,
+        use_mla=False,
+    )
+    assert spec.max_memory_usage_bytes(vllm_config) == 268435456
+    assert spec.page_size_bytes == 262144
+
+    # layer_1 shares KV cache with layer_0
+    spec_shared_0 = copy.copy(spec)
+    spec_shared_0.kv_sharing_target_layer_idx = 0
+    assert spec_shared_0.max_memory_usage_bytes(vllm_config) == 0
+
+    kv_cache_spec: dict[str, KVCacheSpec] = {
+        "layer_0": spec,
+        "layer_1": spec_shared_0,
+    }
+
+    available_memory = 268435456
+
+    kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
+                                          available_memory)
+    assert kv_cache_config.num_blocks == 1024
+    assert kv_cache_config.tensors["layer_0"].size == available_memory
+    assert kv_cache_config.tensors["layer_1"].size == 0
+    assert len(kv_cache_config.kv_cache_groups) == 1
+    kv_sharing_layer_mapping = (
+        kv_cache_config.kv_cache_groups[0].kv_sharing_layer_mapping)
+    assert kv_sharing_layer_mapping is not None
+    assert len(kv_sharing_layer_mapping) == 1
+    assert kv_sharing_layer_mapping['layer_1'] == 0
+
+
 @pytest.mark.parametrize(
-    ("model_id", "max_model_len", "want_estimated_max_len"), [
-        ("Qwen/Qwen1.5-7B", 16385, 16384),
-        ("Qwen/Qwen1.5-7B", 16383, 16383),
-    ])
+    ("model_id", "max_model_len", "want_estimated_max_len",
+     "kv_sharing_factor", "available_mem_gb"), [
+         ("Qwen/Qwen1.5-7B", 16385, 16384, 0, 8),
+         ("Qwen/Qwen1.5-7B", 16383, 16383, 0, 8),
+         ("Qwen/Qwen1.5-7B", 16383, 16383, 2, 4),
+     ])
 def test_estimate_max_model_len(model_id, max_model_len,
-                                want_estimated_max_len):
+                                want_estimated_max_len, kv_sharing_factor,
+                                available_mem_gb):
     # Create a VllmConfig
     model_config = ModelConfig(
         model_id,
@@ -585,17 +650,27 @@ def test_estimate_max_model_len(model_id, max_model_len,
     # Create KV cache specs
     kv_cache_spec = {}
     for i in range(32):
+        kv_sharing_target_layer_idx = None
+        if kv_sharing_factor > 0:
+            share_kv = (i + 1) % kv_sharing_factor == 0
+            if share_kv:
+                # layer idx 1 will use KV cache from idx 0, etc
+                kv_sharing_target_layer_idx = i - (kv_sharing_factor - 1)
+
         layer_name = f"layer_{i}"
-        kv_cache_spec[layer_name] = FullAttentionSpec(
+        spec = FullAttentionSpec(
             block_size=16,
             num_kv_heads=32,
             head_size=128,
             dtype=torch.float16,
             use_mla=False,
         )
-    # Estimate the maximum model length, 16384 model_len need 8GB
+        spec.kv_sharing_target_layer_idx = kv_sharing_target_layer_idx
+        kv_cache_spec[layer_name] = spec
+    # Estimate the maximum model length, 16384 model_len need 8GB normally
+    # with cross-layer KV sharing with sharing factor=2, we only need 4GB
     estimated_max_len = estimate_max_model_len(vllm_config, kv_cache_spec,
-                                               8 * GiB_bytes)
+                                               available_mem_gb * GiB_bytes)
     assert estimated_max_len == want_estimated_max_len
 
 

tests/v1/tpu/worker/test_tpu_model_runner.py

@@ -3,10 +3,13 @@
 
 import pytest
 
+from vllm.attention import Attention
 from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
 from vllm.sampling_params import SamplingParams
 from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
                                        SchedulerOutput)
+from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
+                                        KVCacheGroupSpec, KVCacheTensor)
 from vllm.v1.worker.tpu_model_runner import (
     TPUModelRunner, _get_padded_num_reqs_with_upper_limit,
     _get_padded_token_len, _get_req_paddings, _get_token_paddings)
@@ -292,6 +295,81 @@ def test_update_states_request_unscheduled(model_runner):
     assert not _is_req_scheduled(model_runner, req_ids[1])
 
 
+def test_init_kv_cache_shared_valid(model_runner):
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=model_runner.model_config.get_num_kv_heads(
+            model_runner.parallel_config),
+        head_size=model_runner.model_config.get_head_size(),
+        dtype=model_runner.kv_cache_dtype,
+        use_mla=False,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=10,
+        tensors={
+            "layer.0": KVCacheTensor(size=spec.page_size_bytes * 12),
+            "layer.1": KVCacheTensor(size=0),
+        },
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                # intentionally switch order to check layer names are sorted
+                # such that layers that reuse KV cache from earlier layers
+                # are processed after all layers that allocate KV cache
+                layer_names=["layer.1", "layer.0"],
+                kv_cache_spec=spec,
+                kv_sharing_layer_mapping={"layer.1": 0}),
+        ])
+
+    fwd_context = (
+        model_runner.vllm_config.compilation_config.static_forward_context)
+    # populate forward context before init kv
+    fwd_context['layer.0'] = Attention(32, 128, 0.1)
+    fwd_context['layer.1'] = Attention(32, 128, 0.1)
+
+    model_runner.initialize_kv_cache(kv_cache_config)
+
+    # check memory references of KV caches for layer 0 and 1 are the same
+    assert id(model_runner.kv_caches[0]) == id(model_runner.kv_caches[1])
+    assert len(fwd_context["layer.0"].kv_cache) == 1
+    assert len(fwd_context["layer.1"].kv_cache) == 1
+    layer_1_kv_cache = fwd_context["layer.1"].kv_cache[0]
+    layer_2_kv_cache = fwd_context["layer.1"].kv_cache[0]
+    assert id(layer_1_kv_cache) == id(layer_2_kv_cache)
+
+
+@pytest.mark.parametrize("target_layer_idx", [1, 2])
+def test_init_kv_cache_shared_invalid(model_runner, target_layer_idx):
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=model_runner.model_config.get_num_kv_heads(
+            model_runner.parallel_config),
+        head_size=model_runner.model_config.get_head_size(),
+        dtype=model_runner.kv_cache_dtype,
+        use_mla=False,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=10,
+        tensors={
+            "layer.1": KVCacheTensor(size=spec.page_size_bytes * 12),
+            "layer.0": KVCacheTensor(size=0),
+        },
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                layer_names=["layer.1", "layer.0"],
+                kv_cache_spec=spec,
+                kv_sharing_layer_mapping={"layer.0": target_layer_idx}),
+        ])
+
+    if target_layer_idx >= 2:
+        error_msg = "2 is an invalid layer index!"
+    else:
+        error_msg = ("layer.0 cannot share KV cache with layer.1 which comes"
+                     " after it")
+
+    with pytest.raises(AssertionError, match=error_msg):
+        model_runner.initialize_kv_cache(kv_cache_config)
+
+
 def test_get_paddings():
     # Bucketed padding
     min_token_size, max_token_size, padding_gap = 16, 512, 64

tests/v1/worker/test_gpu_model_runner.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from vllm.attention import Attention
 from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
                          SchedulerConfig, VllmConfig)
 from vllm.sampling_params import SamplingParams
@@ -48,8 +49,7 @@ def initialize_kv_cache(runner: GPUModelRunner):
     runner.initialize_attn_backend(kv_cache_config)
 
 
-@pytest.fixture
-def model_runner():
+def init_model_runner():
     scheduler_config = SchedulerConfig(
         max_num_seqs=10,
         max_num_batched_tokens=512,
@@ -80,6 +80,17 @@ def model_runner():
 
     device = "cuda"
     runner = GPUModelRunner(vllm_config, device)
+    return runner
+
+
+@pytest.fixture(autouse=True)
+def model_runner(request):
+    runner = init_model_runner()
+
+    if 'skipkvinit' in request.keywords:
+        # do not init kv cache for specific tests
+        return runner
+
     initialize_kv_cache(runner)
     return runner
 
@@ -321,3 +332,80 @@ def test_update_states_request_unscheduled(model_runner):
 
     assert _is_req_added(model_runner, req_ids[1])
     assert not _is_req_scheduled(model_runner, req_ids[1])
+
+
+@pytest.mark.skipkvinit
+def test_init_kv_cache_shared_valid(model_runner):
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=model_runner.model_config.get_num_kv_heads(
+            model_runner.parallel_config),
+        head_size=model_runner.model_config.get_head_size(),
+        dtype=model_runner.kv_cache_dtype,
+        use_mla=False,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=10,
+        tensors={
+            "layer.0": KVCacheTensor(size=spec.page_size_bytes * 12),
+            "layer.1": KVCacheTensor(size=0),
+        },
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                # intentionally switch order to check layer names are sorted
+                # such that layers that reuse KV cache from earlier layers
+                # are processed after all layers that allocate KV cache
+                layer_names=["layer.1", "layer.0"],
+                kv_cache_spec=spec,
+                kv_sharing_layer_mapping={"layer.1": 0}),
+        ])
+
+    fwd_context = (
+        model_runner.vllm_config.compilation_config.static_forward_context)
+    # populate forward context before init kv
+    fwd_context['layer.0'] = Attention(32, 128, 0.1)
+    fwd_context['layer.1'] = Attention(32, 128, 0.1)
+
+    model_runner.initialize_kv_cache(kv_cache_config)
+
+    # check memory references of KV caches for layer 0 and 1 are the same
+    assert id(model_runner.kv_caches[0]) == id(model_runner.kv_caches[1])
+    assert len(fwd_context["layer.0"].kv_cache) == 1
+    assert len(fwd_context["layer.1"].kv_cache) == 1
+    layer_1_kv_cache = fwd_context["layer.1"].kv_cache[0]
+    layer_2_kv_cache = fwd_context["layer.1"].kv_cache[0]
+    assert id(layer_1_kv_cache) == id(layer_2_kv_cache)
+
+
+@pytest.mark.skipkvinit
+@pytest.mark.parametrize("target_layer_idx", [1, 2])
+def test_init_kv_cache_shared_invalid(model_runner, target_layer_idx):
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=model_runner.model_config.get_num_kv_heads(
+            model_runner.parallel_config),
+        head_size=model_runner.model_config.get_head_size(),
+        dtype=model_runner.kv_cache_dtype,
+        use_mla=False,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=10,
+        tensors={
+            "layer.1": KVCacheTensor(size=spec.page_size_bytes * 12),
+            "layer.0": KVCacheTensor(size=0),
+        },
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                layer_names=["layer.1", "layer.0"],
+                kv_cache_spec=spec,
+                kv_sharing_layer_mapping={"layer.0": target_layer_idx}),
+        ])
+
+    if target_layer_idx >= 2:
+        error_msg = "2 is an invalid layer index!"
+    else:
+        error_msg = ("layer.0 cannot share KV cache with layer.1 which comes"
+                     " after it")
+
+    with pytest.raises(AssertionError, match=error_msg):
+        model_runner.initialize_kv_cache(kv_cache_config)

vllm/attention/layer.py

@@ -49,6 +49,7 @@ def __init__(
         use_mla: bool = False,
         prefix: str = "",
         attn_type: str = AttentionType.DECODER,
+        kv_sharing_target_layer_idx: Optional[int] = None,
         **extra_impl_args,
     ) -> None:
         """
@@ -102,6 +103,10 @@ def __init__(
         self.head_size = head_size
         self.num_kv_heads = num_kv_heads
         self.sliding_window = sliding_window
+        self.kv_sharing_target_layer_idx = kv_sharing_target_layer_idx
+        if kv_sharing_target_layer_idx is not None:
+            extra_impl_args['kv_sharing_target_layer_idx'] = (
+                kv_sharing_target_layer_idx)
 
         quant_method = quant_config.get_quant_method(
             self, prefix=prefix) if quant_config else None