Fix cp variable naming & Support prefill context parallel with MLA backend.

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

Author: FENP

tests/distributed/test_context_parallel.py

@@ -31,6 +31,7 @@ class ParallelSetup(NamedTuple):
     tp_size: int
     pp_size: int
     dcp_size: int
+    pcp_size: int
     cp_kv_cache_interleave_size: int
     eager_mode: bool
     chunked_prefill: bool
@@ -55,6 +56,7 @@ def detailed(
         tp_base: int = 4,
         pp_base: int = 1,
         dcp_base: int = 1,
+        pcp_base: int = 1,
         cp_kv_cache_interleave_size: int = 1,
         multi_node_only: bool = False,
         runner: RunnerOption = "auto",
@@ -70,7 +72,8 @@ def detailed(
                             ParallelSetup(
                                 tp_size=tp_base,
                                 pp_size=pp_multiplier * pp_base,
-                                dcp_size=int(dcp_multiplier * tp_base),
+                                dcp_size=max(1, int(dcp_multiplier * tp_base)),
+                                pcp_size=pcp_base,
                                 cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
                                 eager_mode=eager_mode_val,
                                 chunked_prefill=chunked_prefill_val,
@@ -116,6 +119,7 @@ def _compare_cp_with_tp(
         tp_size,
         pp_size,
         dcp_size,
+        pcp_size,
         cp_kv_cache_interleave_size,
         eager_mode,
         chunked_prefill,
@@ -196,7 +200,9 @@ def _compare_cp_with_tp(
         str(pp_size),
         "--decode-context-parallel-size",
         str(dcp_size),
-        "--dcp-kv-cache-interleave-size",
+        "--prefill-context-parallel-size",
+        str(pcp_size),
+        "--cp-kv-cache-interleave-size",
         str(cp_kv_cache_interleave_size),
         "--distributed-executor-backend",
         distributed_backend,
@@ -228,6 +234,8 @@ def _compare_cp_with_tp(
         CPTestSettings.detailed(),
         CPTestSettings.detailed(tp_base=2),
         CPTestSettings.detailed(tp_base=2, cp_kv_cache_interleave_size=64),
+        CPTestSettings.detailed(tp_base=1, pcp_base=2),
+        CPTestSettings.detailed(tp_base=1, pcp_base=2, cp_kv_cache_interleave_size=64),
     ],
     "bigcode/gpt_bigcode-santacoder": [
         CPTestSettings.detailed(),

vllm/attention/backends/abstract.py

@@ -295,8 +295,8 @@ class AttentionImpl(ABC, Generic[T]):
     pcp_world_size: int
     pcp_rank: int
 
-    total_cp_world_size: int
-    total_cp_rank: int
+    cp_world_size: int
+    cp_rank: int
 
     def __new__(cls, *args, **kwargs):
         # use __new__ so that all subclasses will call this
@@ -318,11 +318,11 @@ def __new__(cls, *args, **kwargs):
         except AssertionError:
             self.pcp_world_size = 1
             self.pcp_rank = 0
-        self.total_cp_world_size = self.pcp_world_size * self.dcp_world_size
-        self.total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank
+        self.cp_world_size = self.pcp_world_size * self.dcp_world_size
+        self.cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank
 
         self.need_to_return_lse_for_decode = (
-            self.dcp_world_size > 1 and self.can_return_lse_for_decode
+            self.cp_world_size > 1 and self.can_return_lse_for_decode
         )
         return self
 

vllm/config/parallel.py

@@ -237,12 +237,12 @@ class is dynamically inherited by the worker class. This is used to inject
     """
     cp_kv_cache_interleave_size: int = 1
     """Interleave size of kv_cache storage while using DCP or PCP.
-    For `total_cp_rank = pcp_rank * dcp_world_size + dcp_rank`,
-        and `total_cp_world_size = pcp_world_size * dcp_world_size`.
-    store interleave_size tokens on total_cp_rank i,
-    then store next interleave_size tokens on total_cp_rank i+1.
+    For `cp_rank = pcp_rank * dcp_world_size + dcp_rank`,
+        and `cp_world_size = pcp_world_size * dcp_world_size`.
+    store interleave_size tokens on cp_rank i,
+    then store next interleave_size tokens on cp_rank i+1.
     Interleave_size=1: token-level alignment, where token `i` is stored on
-        total_cp_rank `i % total_cp_world_size`.
+        cp_rank `i % cp_world_size`.
     Interleave_size=block_size: block-level alignment, where tokens are
         first populated to the preceding ranks. Tokens are then stored
         in (rank i+1, block j) only after (rank i, block j) is fully occupied.
@@ -312,11 +312,6 @@ def _validate_parallel_config(self) -> Self:
                     "num_redundant_experts."
                 )
 
-        if self.prefill_context_parallel_size > 1:
-            raise ValueError(
-                "Prefill context parallelism is not fully supported. "
-                "Please set prefill_context_parallel_size to 1."
-            )
         return self
 
     @property

vllm/distributed/parallel_state.py

@@ -1085,9 +1085,6 @@ def get_dcp_group() -> GroupCoordinator:
     return _DCP
 
 
-# kept for backward compatibility
-get_context_model_parallel_group = get_dcp_group
-
 _PP: GroupCoordinator | None = None
 
 

vllm/engine/arg_utils.py

@@ -1898,16 +1898,6 @@ def _set_default_chunked_prefill_and_prefix_caching_args(
         default_chunked_prefill = model_config.is_chunked_prefill_supported
         default_prefix_caching = model_config.is_prefix_caching_supported
 
-        if self.prefill_context_parallel_size > 1:
-            default_chunked_prefill = False
-            default_prefix_caching = False
-            logger.warning_once(
-                "--prefill-context-parallel-size > 1 is not compatible with "
-                "chunked prefill and prefix caching now. Chunked prefill "
-                "and prefix caching have been disabled by default.",
-                scope="local",
-            )
-
         if self.enable_chunked_prefill is None:
             self.enable_chunked_prefill = default_chunked_prefill
 

vllm/model_executor/layers/fused_moe/config.py

@@ -894,6 +894,10 @@ def tp_size(self):
     def dp_size(self):
         return self.moe_parallel_config.dp_size
 
+    @property
+    def pcp_size(self):
+        return self.moe_parallel_config.pcp_size
+
     @property
     def ep_size(self):
         return self.moe_parallel_config.ep_size
@@ -906,6 +910,10 @@ def tp_rank(self):
     def dp_rank(self):
         return self.moe_parallel_config.dp_rank
 
+    @property
+    def pcp_rank(self):
+        return self.moe_parallel_config.pcp_rank
+
     @property
     def ep_rank(self):
         return self.moe_parallel_config.ep_rank

vllm/v1/attention/backends/flash_attn.py

@@ -45,7 +45,7 @@
     AttentionCGSupport,
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
-    get_dcp_local_seq_lens,
+    get_cp_local_seq_lens,
     get_kv_cache_layout,
 )
 from vllm.v1.kv_cache_interface import AttentionSpec
@@ -401,7 +401,7 @@ def schedule(
             query_kv_lens = query_start_loc[1:] - query_start_loc[:-1]
             dcp_context_kv_lens = seq_lens - query_kv_lens
 
-            dcp_context_kv_lens = get_dcp_local_seq_lens(
+            dcp_context_kv_lens = get_cp_local_seq_lens(
                 dcp_context_kv_lens,
                 self.dcp_world_size,
                 self.dcp_rank,

vllm/v1/attention/backends/flashinfer.py

@@ -53,7 +53,7 @@
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
     KVCacheLayoutType,
-    get_dcp_local_seq_lens,
+    get_cp_local_seq_lens,
     get_kv_cache_layout,
     get_per_layer_parameters,
     infer_global_hyperparameters,
@@ -694,7 +694,7 @@ def build(
                     seq_lens_cpu[num_decodes:] - query_lens_prefill_cpu
                 )
 
-            seq_lens_cpu = get_dcp_local_seq_lens(
+            seq_lens_cpu = get_cp_local_seq_lens(
                 seq_lens_cpu,
                 self.dcp_world_size,
                 self.dcp_rank,