A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 15:00 UTC
[Misc][PCP&DCP] relocate PCP feature check (#30050)
Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Qiu committed
Dec 11, 2025 at 19:36 UTC
a11f4a81e027efd9ef783b943489c222950ac989
6 files changed
+52
-34
vllm/attention/backends/abstract.py
+6
@@ -294,6 +294,12 @@ class AttentionImpl(ABC, Generic[T]):
294
# Some features like decode context parallelism require the softmax lse.
295
can_return_lse_for_decode: bool = False
296
297
+ # Whether the attention impl supports Prefill Context Parallelism.
298
+ supports_pcp: bool = False
299
+ # Whether the attention impl(or ops) supports MTP
300
+ # when cp_kv_cache_interleave_size > 1
301
+ supports_mtp_with_cp_non_trivial_interleave_size: bool = False
302
+
303
# some attention backends might not always want to return lse
304
# even if they can return lse (for efficiency reasons)
305
need_to_return_lse_for_decode: bool = False
vllm/config/parallel.py
-5
@@ -317,11 +317,6 @@ class ParallelConfig:
317
"num_redundant_experts."
318
)
319
320
- if self.prefill_context_parallel_size > 1:
321
- raise ValueError(
322
- "Prefill context parallelism is not fully supported. "
323
- "Please set prefill_context_parallel_size to 1."
324
- )
320
return self
321
322
@property
vllm/config/vllm.py
-5
@@ -820,11 +820,6 @@ class VllmConfig:
820
f"({self.parallel_config.cp_kv_cache_interleave_size})."
821
)
822
823
- assert (
824
- self.parallel_config.cp_kv_cache_interleave_size == 1
825
- or self.speculative_config is None
826
- ), "MTP with cp_kv_cache_interleave_size > 1 is not supported now."
827
-
823
# Do this after all the updates to compilation_config.mode
824
self.compilation_config.set_splitting_ops_for_v1(
825
all2all_backend=self.parallel_config.all2all_backend,
vllm/engine/arg_utils.py
-10
@@ -1848,16 +1848,6 @@ class EngineArgs:
1848
default_chunked_prefill = model_config.is_chunked_prefill_supported
1849
default_prefix_caching = model_config.is_prefix_caching_supported
1850
1851
- if self.prefill_context_parallel_size > 1:
1852
- default_chunked_prefill = False
1853
- default_prefix_caching = False
1854
- logger.warning_once(
1855
- "--prefill-context-parallel-size > 1 is not compatible with "
1856
- "chunked prefill and prefix caching now. Chunked prefill "
1857
- "and prefix caching have been disabled by default.",
1858
- scope="local",
1859
- )
1860
-
1851
if self.enable_chunked_prefill is None:
1852
self.enable_chunked_prefill = default_chunked_prefill
1853
vllm/v1/worker/cp_utils.py
new
+42
@@ -0,0 +1,42 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import TYPE_CHECKING, Any, cast
4
+
5
+from vllm.config import VllmConfig, get_layers_from_vllm_config
6
+
7
+if TYPE_CHECKING:
8
+ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
9
+else:
10
+ AttentionLayerBase = object
11
+
12
+
13
+def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:
14
+ pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
15
+ dcp_size = vllm_config.parallel_config.decode_context_parallel_size
16
+ interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size
17
+ if pcp_size * dcp_size > 1:
18
+ layer_type = cast(type[Any], AttentionLayerBase)
19
+ layers = get_layers_from_vllm_config(vllm_config, layer_type)
20
+ for layer in layers.values():
21
+ layer_impl = getattr(layer, "impl", None)
22
+ if layer_impl is None:
23
+ continue
24
+ if vllm_config.speculative_config is not None and interleave_size > 1:
25
+ assert layer_impl.supports_mtp_with_cp_non_trivial_interleave_size, (
26
+ "MTP with cp_kv_cache_interleave_size > 1 is not "
27
+ f"supported in {layer_impl.__class__.__name__}."
28
+ )
29
+ if dcp_size > 1:
30
+ assert layer_impl.need_to_return_lse_for_decode, (
31
+ "DCP requires attention impls to return"
32
+ " the softmax lse for decode, but the impl "
33
+ f"{layer_impl.__class__.__name__} "
34
+ "does not return the softmax lse for decode."
35
+ )
36
+
37
+ if pcp_size > 1:
38
+ assert layer_impl.supports_pcp, (
39
+ "PCP requires attention impls' support, "
40
+ f"but the impl {layer_impl.__class__.__name__} "
41
+ "does not support PCP."
42
+ )
vllm/v1/worker/gpu_model_runner.py
+4
-14
@@ -148,6 +148,7 @@ from vllm.v1.spec_decode.ngram_proposer import NgramProposer
148
from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer
149
from vllm.v1.structured_output.utils import apply_grammar_bitmask
150
from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext
151
+from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
152
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
153
from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin
154
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
@@ -4736,6 +4737,9 @@ class GPUModelRunner(
4737
attention_backend_list, kv_cache_config.kv_cache_groups
4738
)
4739
4740
+ # Check if attention backend supports PCP&DCP and related features.
4741
+ check_attention_cp_compatibility(self.vllm_config)
4742
+
4743
for i, attn_backend_map in enumerate(attention_backend_maps):
4744
self.attn_groups.append(create_attn_groups(attn_backend_map, i))
4745
@@ -5394,20 +5398,6 @@ class GPUModelRunner(
5398
kv_transfer_group.register_kv_caches(kv_caches)
5399
kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
5400
5397
- if self.dcp_world_size > 1:
5398
- layer_type = cast(type[Any], AttentionLayerBase)
5399
- layers = get_layers_from_vllm_config(self.vllm_config, layer_type)
5400
- for layer in layers.values():
5401
- layer_impl = getattr(layer, "impl", None)
5402
- if layer_impl is None:
5403
- continue
5404
- assert layer_impl.need_to_return_lse_for_decode, (
5405
- "DCP requires attention impls to return"
5406
- " the softmax lse for decode, but the impl "
5407
- f"{layer_impl.__class__.__name__} "
5408
- "does not return the softmax lse for decode."
5409
- )
5410
-
5401
def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
5402
"""
5403
Add encoder-only layers to the KV cache config.