[DCP] Support dcp kv_cache interleave size > 1

Signed-off-by: zhangsicheng5 <zhangsicheng5@huawei.com>

Author: zhangsicheng5

tests/distributed/test_context_parallel.py

@@ -30,6 +30,7 @@ class ParallelSetup(NamedTuple):
  tp_size: int
  pp_size: int
  dcp_size: int
+ cp_kv_cache_interleave_size: int
  eager_mode: bool
  chunked_prefill: bool
 
@@ -52,6 +53,7 @@ def detailed(
  tp_base: int = 4,
  pp_base: int = 1,
  dcp_base: int = 1,
+ cp_kv_cache_interleave_size: int = 1,
  multi_node_only: bool = False,
  runner: RunnerOption = "auto",
  load_format: str | None = None,
@@ -66,6 +68,7 @@ def detailed(
  tp_size=tp_base,
  pp_size=pp_multiplier * pp_base,
  dcp_size=int(dcp_multiplier * tp_base),
+ cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
  eager_mode=eager_mode_val,
  chunked_prefill=chunked_prefill_val,
  )
@@ -108,6 +111,7 @@ def _compare_cp_with_tp(
  tp_size,
  pp_size,
  dcp_size,
+ cp_kv_cache_interleave_size,
  eager_mode,
  chunked_prefill,
  ) = parallel_setup
@@ -180,6 +184,8 @@ def _compare_cp_with_tp(
  str(pp_size),
  "--decode-context-parallel-size",
  str(dcp_size),
+ "--cp-kv-cache-interleave-size",
+ str(cp_kv_cache_interleave_size),
  "--distributed-executor-backend",
  distributed_backend,
  ]
@@ -208,6 +214,7 @@ def _compare_cp_with_tp(
  "deepseek-ai/DeepSeek-V2-Lite-Chat": [
  CPTestSettings.detailed(),
  CPTestSettings.detailed(tp_base=2),
+ CPTestSettings.detailed(tp_base=2, cp_kv_cache_interleave_size=64),
  ],
 }
 

vllm/config/parallel.py

@@ -204,6 +204,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 = Field(default=1, gt=0)
  """
  The number of API processes initialized.

vllm/config/vllm.py

@@ -472,6 +472,23 @@ def __post_init__(self):
  )
  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})."
+ )
+
+ assert (
+ self.parallel_config.cp_kv_cache_interleave_size == 1
+ or self.speculative_config is None
+ ), "MTP with cp_kv_cache_interleave_size > 1 is not supported now."
+
  # Do this after all the updates to compilation_config.level
  if (
  envs.VLLM_USE_V1

vllm/engine/arg_utils.py

@@ -362,6 +362,7 @@ class EngineArgs:
  pipeline_parallel_size: int = ParallelConfig.pipeline_parallel_size
  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
  data_parallel_size: int = ParallelConfig.data_parallel_size
  data_parallel_rank: int | None = None
  data_parallel_start_rank: int | None = None
@@ -715,6 +716,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
  "-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(
  "--data-parallel-size", "-dp", **parallel_kwargs["data_parallel_size"]
  )
@@ -1470,6 +1475,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/utils/__init__.py

@@ -3426,3 +3426,35 @@ def unique_filepath(fn: Callable[[int], Path]) -> Path:
  if not p.exists():
  return p
  i += 1
+
+
+def get_dcp_local_seq_lens(
+ seq_lens: torch.Tensor,
+ dcp_world_size: int = 1,
+ cp_kv_cache_interleave_size: int = 1,
+) -> torch.Tensor:
+ """While using dcp, kv_cache size stored on each rank may be different,
+ use this function to calculate split decode seq_lens of each dcp rank.
+ Only consider dcp now, we can extend the case of cp based on this.
+ """
+ num_requests = seq_lens.size(0)
+ seq_lens_tiled = seq_lens.unsqueeze(-1).repeat(1, dcp_world_size)
+ rank_offsets = (
+ torch.arange(dcp_world_size, device=seq_lens.device, dtype=torch.int32)
+ .unsqueeze(0)
+ .repeat(num_requests, 1)
+ )
+ base = (
+ seq_lens_tiled
+ // cp_kv_cache_interleave_size
+ // dcp_world_size
+ * cp_kv_cache_interleave_size
+ )
+ remainder = seq_lens_tiled - base * dcp_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
+ return dcp_local_seq_lens

vllm/v1/attention/backends/mla/common.py

@@ -749,15 +749,6 @@ def build(
  )
  )
 
- # Note(hc): update seq_lens of decode reqs under DCP.
- if self.dcp_world_size > 1:
- assert dcp_local_seq_lens is not None
- dcp_local_seq_lens[:num_decodes] = seq_lens[
- :num_decodes
- ] // self.dcp_world_size + (
- self.dcp_rank <= (seq_lens[:num_decodes] - 1) % self.dcp_world_size
- )
-
  assert num_decodes + num_prefills == num_reqs
  assert num_decode_tokens + num_prefill_tokens == num_tokens
 

vllm/v1/worker/block_table.py

@@ -119,7 +119,10 @@ def swap_row(self, src: int, tgt: int) -> None:
  self.block_table.np[src_tgt] = self.block_table.np[tgt_src]
 
  def compute_slot_mapping(
- self, req_indices: np.ndarray, positions: np.ndarray
+ self,
+ req_indices: np.ndarray,
+ positions: np.ndarray,
+ cp_kv_cache_interleave_size: int = 1,
  ) -> None:
  # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
  # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]
@@ -144,9 +147,19 @@ def compute_slot_mapping(
  # Use virtual_block_size for mask calculation, which marks local
  # tokens.
  virtual_block_offsets = positions % virtual_block_size
- mask = virtual_block_offsets % self.dcp_world_size == self.dcp_rank
+ mask = (
+ virtual_block_offsets
+ // cp_kv_cache_interleave_size
+ % self.dcp_world_size
+ == self.dcp_rank
+ )
  # Calculate local block_offsets
- block_offsets = virtual_block_offsets // self.dcp_world_size
+ block_offsets = (
+ virtual_block_offsets
+ // (self.dcp_world_size * cp_kv_cache_interleave_size)
+ * cp_kv_cache_interleave_size
+ + virtual_block_offsets % cp_kv_cache_interleave_size
+ )
  # Calculate slot_mapping
  slot_mapping = block_numbers * self.block_size + block_offsets
  # Write final slots, use -1 for not-local
@@ -284,10 +297,17 @@ def swap_row(self, src: int, tgt: int) -> None:
  block_table.swap_row(src, tgt)
 
  def compute_slot_mapping(
- self, req_indices: np.ndarray, positions: np.ndarray
+ self,
+ req_indices: np.ndarray,
+ positions: np.ndarray,
+ cp_kv_cache_interleave_size: int = 1,
  ) -> None:
  for block_table in self.block_tables:
- block_table.compute_slot_mapping(req_indices, positions)
+ block_table.compute_slot_mapping(
+ req_indices,
+ positions,
+ cp_kv_cache_interleave_size,
+ )
 
  def commit_block_table(self, num_reqs: int) -> None:
  for block_table in self.block_tables:

vllm/v1/worker/gpu_model_runner.py

@@ -35,6 +35,7 @@
 from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
 from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
 from vllm.distributed.parallel_state import (
+ get_dcp_group,
  get_pp_group,
  get_tp_group,
  graph_capture,
@@ -78,6 +79,7 @@
  GiB_bytes,
  cdiv,
  check_use_alibi,
+ get_dcp_local_seq_lens,
  get_dtype_size,
  is_pin_memory_available,
  length_from_prompt_token_ids_or_embeds,
@@ -256,6 +258,11 @@ def __init__(
  self.is_multimodal_pruning_enabled = False
  self.max_model_len = model_config.max_model_len
  self.dcp_world_size = self.parallel_config.decode_context_parallel_size
+ try:
+ self.dcp_rank = get_dcp_group().rank_in_group
+ except AssertionError:
+ # DCP might not be initialized in testing
+ self.dcp_rank = 0
  self.max_num_tokens = scheduler_config.max_num_batched_tokens
  self.max_num_reqs = scheduler_config.max_num_seqs
 
@@ -1158,7 +1165,11 @@ def _prepare_inputs(
 
  output_idx += num_sched
 
- self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np)
+ self.input_batch.block_table.compute_slot_mapping(
+ req_indices,
+ positions_np,
+ self.parallel_config.cp_kv_cache_interleave_size,
+ )
  self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens)
 
  # Prepare the attention metadata.
@@ -1276,6 +1287,14 @@ def _prepare_inputs(
  logits_indices
  )
 
+ # update seq_lens of decode reqs under DCP.
+ if self.dcp_world_size > 1:
+ self.dcp_local_seq_lens.gpu[:num_reqs] = get_dcp_local_seq_lens(
+ seq_lens,
+ self.dcp_world_size,
+ self.parallel_config.cp_kv_cache_interleave_size,
+ )[:, self.dcp_rank]
+
  attn_metadata: PerLayerAttnMetadata = {}
  if ubatch_slices is not None:
  attn_metadata = [dict() for _ in range(len(ubatch_slices))]