Merge pull request #7 from zhangsicheng5/long_seq_dev

support kv_cache interleave_size

Author: LookAround0301

vllm/config/__init__.py

@@ -618,6 +618,18 @@ def __post_init__(self):
                            "to True to enable.")
         current_platform.check_and_update_config(self)
 
+        assert (
+            self.parallel_config.cp_kv_cache_interleave_size
+            <= self.cache_config.block_size
+            and self.cache_config.block_size
+            % self.parallel_config.cp_kv_cache_interleave_size
+            == 0
+        ), (
+            f"Block_size({self.cache_config.block_size}) should be "
+            "greater than or equal to and divisible by cp_kv_cache_interleave_size "
+            f"({self.parallel_config.cp_kv_cache_interleave_size})."
+        )
+
         # final check of cudagraph mode after platform-specific update
         if envs.VLLM_USE_V1 and current_platform.is_cuda_alike():
             if self.compilation_config.cudagraph_mode == CUDAGraphMode.FULL \

vllm/config/parallel.py

@@ -195,6 +195,17 @@ class is dynamically inherited by the worker class. This is used to inject
     not change by dcp, it simply reuse the GPUs of TP group, and tp_size
     needs to be divisible by dcp_size."""
 
+    cp_kv_cache_interleave_size: int = 1
+    """Interleave size of kv_cache storage while using dcp or cp > 1,
+    store interleave_size tokens on (d)cp i,
+    then store next interleave_size tokens on (d)cp i+1.
+    Interleave_size=1: token-level align, token i is stored on rank i % (d)cp_size.
+    Interleave_size=block_size: block-level align, first fill the block on first rank,
+    token is stored on rank i+1 block j after rank i block j is full.
+    Block_size should be greater than or equal to cp_kv_cache_interleave_size.
+    Block_size should be divisible by cp_kv_cache_interleave_size.
+    """
+
     _api_process_count: int = 1
     """
     The number of API processes initialized.

vllm/engine/arg_utils.py

@@ -318,6 +318,7 @@ class EngineArgs:
     tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
     decode_context_parallel_size: int = \
         ParallelConfig.decode_context_parallel_size
+    cp_kv_cache_interleave_size: int = ParallelConfig.cp_kv_cache_interleave_size
     context_parallel_size: int = ParallelConfig.context_parallel_size
     data_parallel_size: int = ParallelConfig.data_parallel_size
     data_parallel_rank: Optional[int] = None
@@ -654,6 +655,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         parallel_group.add_argument(
             "--decode-context-parallel-size", "-dcp",
             **parallel_kwargs["decode_context_parallel_size"])
+        parallel_group.add_argument(
+            "--cp-kv-cache-interleave-size",
+            **parallel_kwargs["cp_kv_cache_interleave_size"])
         parallel_group.add_argument(
             "--context-parallel-size", "-cp",
             **parallel_kwargs["context_parallel_size"])
@@ -1338,6 +1342,7 @@ def create_engine_config(
             worker_cls=self.worker_cls,
             worker_extension_cls=self.worker_extension_cls,
             decode_context_parallel_size=self.decode_context_parallel_size,
+            cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
             _api_process_count=self._api_process_count,
             _api_process_rank=self._api_process_rank,
         )

vllm/v1/attention/backends/utils.py

@@ -851,3 +851,36 @@ def __init__(self, metadata, common_attn_metadata):
         builder_cls=FastPrefillAttentionBuilder)
 
     return attn_backend
+
+
+def get_cp_local_seq_lens(
+    seq_lens: torch.Tensor,
+    cp_world_size: int = 1,
+    dcp_world_size: int = 1,
+    cp_kv_cache_interleave_size: int = 1,
+) -> torch.Tensor:
+    """While using cp or dcp, kv_cache size stored on each rank may be different,
+    use this function to calculate split decode seq_lens of each (d)cp rank.
+    """
+    num_requests = seq_lens.size(0)
+    total_world_size = cp_world_size * dcp_world_size
+    seq_lens_tiled = seq_lens.unsqueeze(-1).repeat(1, total_world_size)
+    rank_offsets = (
+        torch.arange(total_world_size, dtype=torch.int32)
+        .unsqueeze(0)
+        .repeat(num_requests, 1)
+    )
+    base = (
+        seq_lens_tiled
+        // cp_kv_cache_interleave_size
+        // total_world_size
+        * cp_kv_cache_interleave_size
+    )
+    remainder = seq_lens_tiled - base * total_world_size
+    remainder = torch.clip(
+        remainder - rank_offsets * cp_kv_cache_interleave_size,
+        0,
+        cp_kv_cache_interleave_size,
+    )
+    dcp_local_seq_lens = (base + remainder).reshape([-1, cp_world_size, dcp_world_size])
+    return dcp_local_seq_lens