@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2fd893b4ce
Fork 0

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 12, 22:15 UTC

[Feature] Prefill Context Parallel (PCP) basic support (#28718)

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com> Signed-off-by: FENP <yuanyongjie.yyj@antgroup.com> Signed-off-by: LookAround <lixushi@huawei.com> Signed-off-by: Jingchun Gao <gaojingchun1@huawei.com> Signed-off-by: zhenwenqi2024 <zhenwenqi_2022@qq.com> Co-authored-by: FENP <yuanyongjie.yyj@antgroup.com> Co-authored-by: LookAround <lixushi@huawei.com> Co-authored-by: Jingchun Gao <gaojingchun1@huawei.com> Co-authored-by: zhenwenqi2024 <zhenwenqi_2022@qq.com> Co-authored-by: Jingchun Gao <63247409+gjc0824@users.noreply.github.com>

Qiu committed Nov 20, 2025 at 04:52 UTC 2fd893b4cec0975a2a8430077fd9b4f294eb3561
27 files changed +399 -114
tests/distributed/test_context_parallel.py
+6 -6
@@ -31,7 +31,7 @@ class ParallelSetup(NamedTuple):
31 tp_size: int
32 pp_size: int
33 dcp_size: int
34 - dcp_kv_cache_interleave_size: int
34 + cp_kv_cache_interleave_size: int
35 eager_mode: bool
36 chunked_prefill: bool
37
@@ -55,7 +55,7 @@ class CPTestSettings:
55 tp_base: int = 4,
56 pp_base: int = 1,
57 dcp_base: int = 1,
58 - dcp_kv_cache_interleave_size: int = 1,
58 + cp_kv_cache_interleave_size: int = 1,
59 multi_node_only: bool = False,
60 runner: RunnerOption = "auto",
61 load_format: str | None = None,
@@ -71,7 +71,7 @@ class CPTestSettings:
71 tp_size=tp_base,
72 pp_size=pp_multiplier * pp_base,
73 dcp_size=int(dcp_multiplier * tp_base),
74 - dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,
74 + cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
75 eager_mode=eager_mode_val,
76 chunked_prefill=chunked_prefill_val,
77 )
@@ -116,7 +116,7 @@ def _compare_cp_with_tp(
116 tp_size,
117 pp_size,
118 dcp_size,
119 - dcp_kv_cache_interleave_size,
119 + cp_kv_cache_interleave_size,
120 eager_mode,
121 chunked_prefill,
122 ) = parallel_setup
@@ -197,7 +197,7 @@ def _compare_cp_with_tp(
197 "--decode-context-parallel-size",
198 str(dcp_size),
199 "--dcp-kv-cache-interleave-size",
200 - str(dcp_kv_cache_interleave_size),
200 + str(cp_kv_cache_interleave_size),
201 "--distributed-executor-backend",
202 distributed_backend,
203 ]
@@ -227,7 +227,7 @@ CP_TEXT_GENERATION_MODELS = {
227 "deepseek-ai/DeepSeek-V2-Lite-Chat": [
228 CPTestSettings.detailed(),
229 CPTestSettings.detailed(tp_base=2),
230 - CPTestSettings.detailed(tp_base=2, dcp_kv_cache_interleave_size=64),
230 + CPTestSettings.detailed(tp_base=2, cp_kv_cache_interleave_size=64),
231 ],
232 "bigcode/gpt_bigcode-santacoder": [
233 CPTestSettings.detailed(),
tests/kernels/moe/modular_kernel_tools/common.py
+6 -1
@@ -15,7 +15,11 @@ from tests.kernels.quantization.nvfp4_utils import (
15 )
16 from tests.kernels.utils import torch_experts
17 from vllm.config import VllmConfig
18 -from vllm.distributed import get_dp_group, get_tensor_model_parallel_world_size
18 +from vllm.distributed import (
19 + get_dp_group,
20 + get_pcp_group,
21 + get_tensor_model_parallel_world_size,
22 +)
23 from vllm.forward_context import set_forward_context
24 from vllm.model_executor.layers.fused_moe.config import (
25 FusedMoEConfig,
@@ -561,6 +565,7 @@ def make_modular_kernel(
565 # make moe config
566 moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(
567 tp_size_=get_tensor_model_parallel_world_size(),
568 + pcp_size_=get_pcp_group().world_size,
569 dp_size_=get_dp_group().world_size,
570 vllm_parallel_config=vllm_config.parallel_config,
571 )
tests/v1/worker/test_gpu_model_runner.py
+2 -2
@@ -956,7 +956,7 @@ def test_hybrid_block_table_initialization():
956 max_num_reqs = 10
957 max_num_blocks_per_req = 20
958 max_num_batched_tokens = 512
959 - dcp_kv_cache_interleave_size = 8
959 + cp_kv_cache_interleave_size = 8
960
961 block_table = BlockTable(
962 block_size=block_size,
@@ -966,7 +966,7 @@ def test_hybrid_block_table_initialization():
966 pin_memory=False,
967 device=torch.device(DEVICE),
968 kernel_block_size=kernel_block_sizes[0],
969 - dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,
969 + cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
970 )
971
972 # Verify hybrid block configuration
vllm/attention/backends/abstract.py
+17
@@ -266,6 +266,12 @@ class AttentionImpl(ABC, Generic[T]):
266 dcp_world_size: int
267 dcp_rank: int
268
269 + pcp_world_size: int
270 + pcp_rank: int
271 +
272 + total_cp_world_size: int
273 + total_cp_rank: int
274 +
275 def __new__(cls, *args, **kwargs):
276 # use __new__ so that all subclasses will call this
277 self = super().__new__(cls)
@@ -278,6 +284,17 @@ class AttentionImpl(ABC, Generic[T]):
284 # DCP might not be initialized in testing
285 self.dcp_world_size = 1
286 self.dcp_rank = 0
287 + try:
288 + from vllm.distributed.parallel_state import get_pcp_group
289 +
290 + self.pcp_world_size = get_pcp_group().world_size
291 + self.pcp_rank = get_pcp_group().rank_in_group
292 + except AssertionError:
293 + self.pcp_world_size = 1
294 + self.pcp_rank = 0
295 + self.total_cp_world_size = self.pcp_world_size * self.dcp_world_size
296 + self.total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank
297 +
298 self.need_to_return_lse_for_decode = (
299 self.dcp_world_size > 1 and self.can_return_lse_for_decode
300 )
vllm/attention/ops/common.py
+37 -3
@@ -169,12 +169,11 @@ def correct_attn_out(
169 return out, lse
170
171
172 -def cp_lse_ag_out_rs(
172 +def _cp_lse_common(
173 cp_attn_out: torch.Tensor,
174 cp_attn_lse: torch.Tensor,
175 cp_group: GroupCoordinator,
176 - ctx: CPTritonContext = None,
177 - return_lse=False,
176 + ctx: CPTritonContext | None = None,
177 ):
178 """
179 cp_attn_out: [ B, H, D ]
@@ -195,6 +194,22 @@ def cp_lse_ag_out_rs(
194 cp_attn_lse = cp_attn_lse.contiguous()
195 lses = cp_group.all_gather(cp_attn_lse, dim=0).view_as(lses)
196 out, lse = correct_attn_out(cp_attn_out, lses, cp_group.rank_in_group, ctx)
197 + assert out.is_contiguous()
198 + return out, lse
199 +
200 +
201 +def cp_lse_ag_out_rs(
202 + cp_attn_out: torch.Tensor,
203 + cp_attn_lse: torch.Tensor,
204 + cp_group: GroupCoordinator,
205 + ctx: CPTritonContext | None = None,
206 + return_lse: bool = False,
207 +):
208 + """
209 + cp_attn_out: [ B, H, D ]
210 + cp_attn_lse: [ B, H ]
211 + """
212 + out, lse = _cp_lse_common(cp_attn_out, cp_attn_lse, cp_group, ctx=ctx)
213 out = cp_group.reduce_scatter(out, dim=1)
214
215 if return_lse:
@@ -205,6 +220,25 @@ def cp_lse_ag_out_rs(
220 return out
221
222
223 +def cp_lse_ag_out_ar(
224 + cp_attn_out: torch.Tensor,
225 + cp_attn_lse: torch.Tensor,
226 + cp_group: GroupCoordinator,
227 + ctx: CPTritonContext | None = None,
228 + return_lse: bool = False,
229 +):
230 + """
231 + cp_attn_out: [ B, H, D ]
232 + cp_attn_lse: [ B, H ]
233 + """
234 + out, lse = _cp_lse_common(cp_attn_out, cp_attn_lse, cp_group, ctx=ctx)
235 + out = cp_group.all_reduce(out)
236 +
237 + if return_lse:
238 + return out, lse
239 + return out
240 +
241 +
242 @triton.jit
243 def _pack_seq_kernel(
244 x_ptr, # [N, D]
vllm/config/parallel.py
+31 -9
@@ -71,6 +71,8 @@ class ParallelConfig:
71 """Number of pipeline parallel groups."""
72 tensor_parallel_size: int = 1
73 """Number of tensor parallel groups."""
74 + prefill_context_parallel_size: int = 1
75 + """Number of prefill context parallel groups."""
76 data_parallel_size: int = 1
77 """Number of data parallel groups. MoE layers will be sharded according to
78 the product of the tensor parallel size and data parallel size."""
@@ -239,14 +241,25 @@ class ParallelConfig:
241 needs to be divisible by dcp_size."""
242
243 dcp_kv_cache_interleave_size: int = 1
242 - """Interleave size of kv_cache storage while using dcp or cp > 1,
243 - store interleave_size tokens on (d)cp i,
244 - then store next interleave_size tokens on (d)cp i+1.
245 - Interleave_size=1: token-level align, token i is stored on rank i % (d)cp_size.
246 - Interleave_size=block_size: block-level align, first fill the block on first rank,
247 - token is stored on rank i+1 block j after rank i block j is full.
248 - Block_size should be greater than or equal to dcp_kv_cache_interleave_size.
249 - Block_size should be divisible by dcp_kv_cache_interleave_size.
244 + """
245 + Interleave size of kv_cache storage while using DCP.
246 + dcp_kv_cache_interleave_size has been replaced by cp_kv_cache_interleave_size,
247 + and will be deprecated when PCP is fully supported.
248 +
249 + """
250 + cp_kv_cache_interleave_size: int = 1
251 + """Interleave size of kv_cache storage while using DCP or PCP.
252 + For `total_cp_rank = pcp_rank * dcp_world_size + dcp_rank`,
253 + and `total_cp_world_size = pcp_world_size * dcp_world_szie`.
254 + store interleave_size tokens on total_cp_rank i,
255 + then store next interleave_size tokens on taotal_cp_rank i+1.
256 + Interleave_size=1: token-level alignment, where token `i` is stored on
257 + total_cp_rank `i % total_cp_world_size`.
258 + Interleave_size=block_size: block-level alignment, where tokens are
259 + first populated to the preceding ranks. Tokens are then stored
260 + in (rank i+1, block j) only after (rank i, block j) is fully occupied.
261 + Block_size should be greater than or equal to cp_kv_cache_interleave_size.
262 + Block_size should be divisible by cp_kv_cache_interleave_size.
263 """
264
265 _api_process_count: int = Field(default=1, gt=0)
@@ -311,6 +324,11 @@ class ParallelConfig:
324 "num_redundant_experts."
325 )
326
327 + if self.prefill_context_parallel_size > 1:
328 + raise ValueError(
329 + "Prefill context parallelism is not fully supported. "
330 + "Please set prefill_context_parallel_size to 1."
331 + )
332 return self
333
334 @property
@@ -529,7 +547,11 @@ class ParallelConfig:
547 )
548
549 # Continue with the rest of the initialization
532 - self.world_size = self.pipeline_parallel_size * self.tensor_parallel_size
550 + self.world_size = (
551 + self.pipeline_parallel_size
552 + * self.tensor_parallel_size
553 + * self.prefill_context_parallel_size
554 + )
555
556 if self.distributed_executor_backend == "external_launcher":
557 logger.info("Using external launcher for distributed inference.")
vllm/config/vllm.py
+26 -6
@@ -481,6 +481,14 @@ class VllmConfig:
481 "Overriding cudagraph_mode to PIECEWISE."
482 )
483 self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
484 + # prefill context parallel do not support full cudagraphs
485 + elif self.parallel_config.prefill_context_parallel_size > 1:
486 + logger.warning_once(
487 + "Prefill context parallel (PCP) is enabled, which is "
488 + "incompatible with full CUDA graphs. "
489 + "Overriding cudagraph_mode to PIECEWISE."
490 + )
491 + self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
492 elif self.model_config is not None:
493 if self.model_config.pooler_config is not None:
494 logger.warning_once(
@@ -610,22 +618,34 @@ class VllmConfig:
618
619 # If DCP, ensure the block size is right.
620 if self.parallel_config.decode_context_parallel_size > 1:
621 + if self.parallel_config.dcp_kv_cache_interleave_size > 1 and (
622 + self.parallel_config.cp_kv_cache_interleave_size
623 + != self.parallel_config.dcp_kv_cache_interleave_size
624 + ):
625 + self.parallel_config.cp_kv_cache_interleave_size = (
626 + self.parallel_config.dcp_kv_cache_interleave_size
627 + )
628 + logger.warning_once(
629 + "cp_kv_cache_interleave_size is overridden by dcp_kv_cache"
630 + "_interleave_size. And dcp-kv-cache-interleave-size will be "
631 + "deprecated when PCP is fully supported."
632 + )
633 assert (
614 - self.parallel_config.dcp_kv_cache_interleave_size
634 + self.parallel_config.cp_kv_cache_interleave_size
635 <= self.cache_config.block_size
636 and self.cache_config.block_size
617 - % self.parallel_config.dcp_kv_cache_interleave_size
637 + % self.parallel_config.cp_kv_cache_interleave_size
638 == 0
639 ), (
640 f"Block_size({self.cache_config.block_size}) should be greater "
621 - "than or equal to and divisible by dcp_kv_cache_interleave_size "
622 - f"({self.parallel_config.dcp_kv_cache_interleave_size})."
641 + "than or equal to and divisible by cp_kv_cache_interleave_size "
642 + f"({self.parallel_config.cp_kv_cache_interleave_size})."
643 )
644
645 assert (
626 - self.parallel_config.dcp_kv_cache_interleave_size == 1
646 + self.parallel_config.cp_kv_cache_interleave_size == 1
647 or self.speculative_config is None
628 - ), "MTP with dcp_kv_cache_interleave_size > 1 is not supported now."
648 + ), "MTP with cp_kv_cache_interleave_size > 1 is not supported now."
649
650 # Do this after all the updates to compilation_config.mode
651 if self.compilation_config.mode == CompilationMode.VLLM_COMPILE:
vllm/distributed/parallel_state.py
+61 -13
@@ -1098,6 +1098,12 @@ get_context_model_parallel_group = get_dcp_group
1098
1099 _PP: GroupCoordinator | None = None
1100
1101 +
1102 +def get_pp_group() -> GroupCoordinator:
1103 + assert _PP is not None, "pipeline model parallel group is not initialized"
1104 + return _PP
1105 +
1106 +
1107 _DP: GroupCoordinator | None = None
1108
1109
@@ -1114,9 +1120,12 @@ def get_ep_group() -> GroupCoordinator:
1120 return _EP
1121
1122
1117 -def get_pp_group() -> GroupCoordinator:
1118 - assert _PP is not None, "pipeline model parallel group is not initialized"
1119 - return _PP
1123 +_PCP: GroupCoordinator | None = None
1124 +
1125 +
1126 +def get_pcp_group() -> GroupCoordinator:
1127 + assert _PCP is not None, "prefill context parallel group is not initialized"
1128 + return _PCP
1129
1130
1131 @deprecated(
@@ -1276,6 +1285,7 @@ def init_distributed_environment(
1285 def initialize_model_parallel(
1286 tensor_model_parallel_size: int = 1,
1287 pipeline_model_parallel_size: int = 1,
1288 + prefill_context_model_parallel_size: int = 1,
1289 decode_context_model_parallel_size: int | None = 1,
1290 backend: str | None = None,
1291 ) -> None:
@@ -1325,7 +1335,11 @@ def initialize_model_parallel(
1335 # to get group_ranks for each dimension, transpose that dimension to the
1336 # last dimension, then reshape to 2D, then unbind the last dimension
1337 all_ranks = torch.arange(world_size).reshape(
1328 - -1, data_parallel_size, pipeline_model_parallel_size, tensor_model_parallel_size
1338 + -1,
1339 + data_parallel_size,
1340 + pipeline_model_parallel_size,
1341 + prefill_context_model_parallel_size,
1342 + tensor_model_parallel_size,
1343 ) # noqa
1344
1345 # Build the tensor model-parallel groups.
@@ -1360,11 +1374,23 @@ def initialize_model_parallel(
1374 group_name="dcp",
1375 )
1376
1377 + global _PCP
1378 + assert _PCP is None, "prefill context parallel group is already initialized"
1379 + group_ranks = (
1380 + all_ranks.transpose(3, 4)
1381 + .reshape(-1, prefill_context_model_parallel_size)
1382 + .unbind(0)
1383 + )
1384 + group_ranks = [x.tolist() for x in group_ranks]
1385 + _PCP = init_model_parallel_group(
1386 + group_ranks, get_world_group().local_rank, backend, group_name="pcp"
1387 + )
1388 +
1389 # Build the pipeline model-parallel groups.
1390 global _PP
1391 assert _PP is None, "pipeline model parallel group is already initialized"
1392 group_ranks = (
1367 - all_ranks.transpose(2, 3).reshape(-1, pipeline_model_parallel_size).unbind(0)
1393 + all_ranks.transpose(2, 4).reshape(-1, pipeline_model_parallel_size).unbind(0)
1394 )
1395 group_ranks = [x.tolist() for x in group_ranks]
1396 _PP = init_model_parallel_group(
@@ -1373,7 +1399,7 @@ def initialize_model_parallel(
1399
1400 global _DP
1401 assert _DP is None, "data parallel group is already initialized"
1376 - group_ranks = all_ranks.transpose(1, 3).reshape(-1, data_parallel_size).unbind(0)
1402 + group_ranks = all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0)
1403 group_ranks = [x.tolist() for x in group_ranks]
1404 _DP = init_model_parallel_group(
1405 group_ranks, get_world_group().local_rank, backend, group_name="dp"
@@ -1383,7 +1409,12 @@ def initialize_model_parallel(
1409 assert _EP is None, "expert parallel group is already initialized"
1410 group_ranks = (
1411 all_ranks.transpose(1, 2)
1386 - .reshape(-1, data_parallel_size * tensor_model_parallel_size)
1412 + .reshape(
1413 + -1,
1414 + data_parallel_size
1415 + * prefill_context_model_parallel_size
1416 + * tensor_model_parallel_size,
1417 + )
1418 .unbind(0)
1419 )
1420 group_ranks = [x.tolist() for x in group_ranks]
@@ -1393,11 +1424,13 @@ def initialize_model_parallel(
1424
1425 logger.info_once(
1426 "rank %s in world size %s is assigned as "
1396 - "DP rank %s, PP rank %s, TP rank %s, EP rank %s",
1427 + "DP rank %s, PP rank %s, PCP rank %s, "
1428 + "TP rank %s, EP rank %s",
1429 rank,
1430 world_size,
1431 _DP.rank_in_group,
1432 _PP.rank_in_group,
1433 + _PCP.rank_in_group,
1434 _TP.rank_in_group,
1435 _EP.rank_in_group,
1436 )
@@ -1406,6 +1439,7 @@ def initialize_model_parallel(
1439 def ensure_model_parallel_initialized(
1440 tensor_model_parallel_size: int,
1441 pipeline_model_parallel_size: int,
1442 + prefill_context_model_parallel_size: int = 1,
1443 decode_context_model_parallel_size: int | None = 1,
1444 backend: str | None = None,
1445 ) -> None:
@@ -1418,6 +1452,7 @@ def ensure_model_parallel_initialized(
1452 initialize_model_parallel(
1453 tensor_model_parallel_size,
1454 pipeline_model_parallel_size,
1455 + prefill_context_model_parallel_size,
1456 decode_context_model_parallel_size,
1457 backend,
1458 )
@@ -1434,6 +1469,12 @@ def ensure_model_parallel_initialized(
1469 f"got: {pp_world_size=} vs. "
1470 f"wanted: {pipeline_model_parallel_size=}"
1471 )
1472 + pcp_world_size = get_pcp_group().world_size
1473 + assert pcp_world_size == prefill_context_model_parallel_size, (
1474 + "prefill context parallel group already initialized, but of unexpected size: "
1475 + f"{pcp_world_size=} vs. "
1476 + f"{prefill_context_model_parallel_size=}"
1477 + )
1478
1479
1480 def prepare_communication_buffer_for_model(model: torch.nn.Module):
@@ -1445,6 +1486,8 @@ def prepare_communication_buffer_for_model(model: torch.nn.Module):
1486 """
1487 if _TP is not None:
1488 _TP.prepare_communication_buffer_for_model(model)
1489 + if _PCP is not None:
1490 + _PCP.prepare_communication_buffer_for_model(model)
1491 if _PP is not None:
1492 _PP.prepare_communication_buffer_for_model(model)
1493 if _DP is not None:
@@ -1520,16 +1563,21 @@ def destroy_model_parallel():
1563 _TP.destroy()
1564 _TP = None
1565
1523 - global _PP
1524 - if _PP:
1525 - _PP.destroy()
1526 - _PP = None
1527 -
1566 global _DCP
1567 if _DCP:
1568 _DCP.destroy()
1569 _DCP = None
1570
1571 + global _PCP
1572 + if _PCP:
1573 + _PCP.destroy()
1574 + _PCP = None
1575 +
1576 + global _PP
1577 + if _PP:
1578 + _PP.destroy()
1579 + _PP = None
1580 +
1581 global _DP
1582 if _DP:
1583 _DP.destroy()
vllm/engine/arg_utils.py
+22
@@ -389,8 +389,10 @@ class EngineArgs:
389 nnodes: int = ParallelConfig.nnodes
390 node_rank: int = ParallelConfig.node_rank
391 tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
392 + prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size
393 decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size
394 dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size
395 + cp_kv_cache_interleave_size: int = ParallelConfig.cp_kv_cache_interleave_size
396 data_parallel_size: int = ParallelConfig.data_parallel_size
397 data_parallel_rank: int | None = None
398 data_parallel_start_rank: int | None = None
@@ -770,6 +772,15 @@ class EngineArgs:
772 "--dcp-kv-cache-interleave-size",
773 **parallel_kwargs["dcp_kv_cache_interleave_size"],
774 )
775 + parallel_group.add_argument(
776 + "--cp-kv-cache-interleave-size",
777 + **parallel_kwargs["cp_kv_cache_interleave_size"],
778 + )
779 + parallel_group.add_argument(
780 + "--prefill-context-parallel-size",
781 + "-pcp",
782 + **parallel_kwargs["prefill_context_parallel_size"],
783 + )
784 parallel_group.add_argument(
785 "--data-parallel-size", "-dp", **parallel_kwargs["data_parallel_size"]
786 )
@@ -1600,6 +1611,7 @@ class EngineArgs:
1611 parallel_config = ParallelConfig(
1612 pipeline_parallel_size=self.pipeline_parallel_size,
1613 tensor_parallel_size=self.tensor_parallel_size,
1614 + prefill_context_parallel_size=self.prefill_context_parallel_size,
1615 data_parallel_size=self.data_parallel_size,
1616 data_parallel_rank=self.data_parallel_rank or 0,
1617 data_parallel_external_lb=data_parallel_external_lb,
@@ -1631,6 +1643,7 @@ class EngineArgs:
1643 worker_extension_cls=self.worker_extension_cls,
1644 decode_context_parallel_size=self.decode_context_parallel_size,
1645 dcp_kv_cache_interleave_size=self.dcp_kv_cache_interleave_size,
1646 + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
1647 _api_process_count=self._api_process_count,
1648 _api_process_rank=self._api_process_rank,
1649 )
@@ -1952,6 +1965,15 @@ class EngineArgs:
1965 default_prefix_caching,
1966 ) = self.get_chunked_prefill_prefix_caching_defaults(model_config)
1967
1968 + if self.prefill_context_parallel_size > 1:
1969 + default_chunked_prefill = False
1970 + default_prefix_caching = False
1971 + logger.warning(
1972 + "--prefill-context-parallel-size > 1 is not compatible with "
1973 + "chunked prefill and prefix caching now. Chunked prefill "
1974 + "and prefix caching have been disabled by default."
1975 + )
1976 +
1977 if self.enable_chunked_prefill is None:
1978 self.enable_chunked_prefill = default_chunked_prefill
1979
vllm/model_executor/layers/fused_moe/config.py
+40 -19
@@ -8,7 +8,11 @@ import torch
8
9 import vllm.envs as envs
10 from vllm.config import ParallelConfig
11 -from vllm.distributed import get_dp_group, get_tensor_model_parallel_rank
11 +from vllm.distributed import (
12 + get_dp_group,
13 + get_pcp_group,
14 + get_tensor_model_parallel_rank,
15 +)
16 from vllm.logger import init_logger
17 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
18 OCP_MX_DTYPES,
@@ -684,9 +688,11 @@ FUSED_MOE_UNQUANTIZED_CONFIG: FusedMoEQuantConfig = FusedMoEQuantConfig.make()
688 @dataclass
689 class FusedMoEParallelConfig:
690 tp_size: int
691 + pcp_size: int
692 dp_size: int
693 ep_size: int
694 tp_rank: int
695 + pcp_rank: int
696 dp_rank: int
697 ep_rank: int
698
@@ -713,19 +719,22 @@ class FusedMoEParallelConfig:
719 return self.use_all2all_kernels and self.all2all_backend == "deepep_low_latency"
720
721 @staticmethod
716 - def flatten_tp_across_dp(
717 - tp_size: int, dp_size: int, dp_rank: int
722 + def flatten_tp_across_dp_and_pcp(
723 + tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int
724 ) -> tuple[int, int]:
725 tp_rank = 0 if tp_size == 1 else get_tensor_model_parallel_rank()
720 - # There are actually dp_size * tp_size devices. Update tp_size
721 - # and tp_rank so we shard across all devices.
722 - flatten_tp_size = dp_size * tp_size
723 - flatten_tp_rank = dp_rank * tp_size + tp_rank
726 + # There are actually dp_size * pcp_size * tp_size devices.
727 + # Update tp_size and tp_rank so we shard across all devices.
728 + flatten_tp_size = dp_size * pcp_size * tp_size
729 + flatten_tp_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank
730 return flatten_tp_size, flatten_tp_rank
731
732 @staticmethod
733 def make(
728 - tp_size_: int, dp_size_: int, vllm_parallel_config: ParallelConfig
734 + tp_size_: int,
735 + pcp_size_: int,
736 + dp_size_: int,
737 + vllm_parallel_config: ParallelConfig,
738 ) -> "FusedMoEParallelConfig":
739 """
740 Determine MoE parallel configuration. Based on the input `tp_size_`,
@@ -734,19 +743,22 @@ class FusedMoEParallelConfig:
743
744 Args:
745 tp_size_ (int): `tp_size` passed into the FusedMoE constructor.
746 + pcp_size_ (int): `pcp_size` passed into the FusedMoE constructor.
747 dp_size_ (int): `dp_size` passed into the FusedMoE constructor.
748 vllm_parallel_config (ParallelConfig): vLLM's parallel config
749 object which contains the `enable_expert_parallel` flag.
750
751 Examples:
752 When there is no parallelism requested,
743 - i.e. `tp_size_` = `dp_size_` = 1, we simply return the sizes
753 + i.e. `tp_size_` = `pcp_size_` = `dp_size_` = 1, we simply return the sizes
754 unaltered and the ranks set to 0.
755
746 - Expert Parallelism is considered only when either `dp_size_` or
756 + Expert Parallelism is considered only when either `dp_size_`, `pcp_size_` or
757 `tp_size_` is non trivial.
758
749 - When TP = 2, DP = 1 and EP = False, the configuration on different
759 + Note that PCP serves the same function as DP here.
760 +
761 + When TP = 2, DP(PCP) = 1 and EP = False, the configuration on different
762 devices:
763
764 - device 0 : TP = {2, 0} DP = {1, 0} EP = {1, 0} //
@@ -754,7 +766,7 @@ class FusedMoEParallelConfig:
766 - device 1 : TP = {2, 1} DP = {1, 0} EP = {1, 0}
767 - Comment : Tensors are sharded across 2 devices.
768
757 - When TP = 1, DP = 2 and EP = False, the configuration on different
769 + When TP = 1, DP(PCP) = 2 and EP = False, the configuration on different
770 devices:
771
772 - device 0 : TP = {2, 0} DP = {2, 0} EP = {1, 0}
@@ -762,7 +774,7 @@ class FusedMoEParallelConfig:
774 - Comment: There are 2 engine instances and the tensors are sharded
775 across 2 decvices.
776
765 - When TP = 2, DP = 2 and EP = False, the configuration on different
777 + When TP = 2, DP(PCP) = 2 and EP = False, the configuration on different
778 devices:
779
780 - device 0: TP = {4, 0} DP = {2, 0} EP = {1, 0}
@@ -772,14 +784,14 @@ class FusedMoEParallelConfig:
784 - Comment: There are 2 engine instances and the tensors are sharded
785 across 4 devices.
786
775 - When, TP = 2, DP = 1 and EP = True, the configuration on different
787 + When, TP = 2, DP(PCP) = 1 and EP = True, the configuration on different
788 devices:
789
790 - device 0: TP = {1, 0} DP = {1, 0} EP = {2, 0}
791 - device 1: TP = {1, 0} DP = {1, 0} EP = {2, 1}
792 - Comment: The experts are split between the 2 devices.
793
782 - When, TP = 1, DP = 2 and EP = True, the configuration on different
794 + When, TP = 1, DP(PCP) = 2 and EP = True, the configuration on different
795 devices:
796
797 - device 0: TP = {1, 0} DP = {2, 0} EP = {2, 0}
@@ -787,7 +799,7 @@ class FusedMoEParallelConfig:
799 - Comment: There are 2 engine instances and the experts are split
800 between the 2 devices.
801
790 - When TP = 2, DP = 2 and EP = True, the configuration on different
802 + When TP = 2, DP(PCP) = 2 and EP = True, the configuration on different
803 devices:
804
805 - device 0: TP = {1, 0} DP = {2, 0} EP = {4, 0}
@@ -798,18 +810,25 @@ class FusedMoEParallelConfig:
810 between the 4 devices.
811 """
812
801 - use_ep = dp_size_ * tp_size_ > 1 and vllm_parallel_config.enable_expert_parallel
813 + use_ep = (
814 + dp_size_ * pcp_size_ * tp_size_ > 1
815 + and vllm_parallel_config.enable_expert_parallel
816 + )
817
818 dp_size = dp_size_
819 dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
805 - tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp(
806 - tp_size_, dp_size_, dp_rank
820 + pcp_size = pcp_size_
821 + pcp_rank = get_pcp_group().rank_in_group if pcp_size > 1 else 0
822 + tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
823 + tp_size_, dp_size_, dp_rank, pcp_size_, pcp_rank
824 )
825
826 if not use_ep:
827 return FusedMoEParallelConfig(
828 tp_size=tp_size,
829 tp_rank=tp_rank,
830 + pcp_size=pcp_size,
831 + pcp_rank=pcp_rank,
832 dp_size=dp_size,
833 dp_rank=dp_rank,
834 ep_size=1,
@@ -826,6 +845,8 @@ class FusedMoEParallelConfig:
845 return FusedMoEParallelConfig(
846 tp_size=1,
847 tp_rank=0,
848 + pcp_size=pcp_size,
849 + pcp_rank=pcp_rank,
850 dp_size=dp_size,
851 dp_rank=dp_rank,
852 ep_size=ep_size,
vllm/model_executor/layers/fused_moe/layer.py
+32
@@ -18,6 +18,7 @@ from vllm.config.parallel import ExpertPlacementStrategy
18 from vllm.distributed import (
19 get_dp_group,
20 get_ep_group,
21 + get_pcp_group,
22 get_tensor_model_parallel_world_size,
23 tensor_model_parallel_all_reduce,
24 )
@@ -343,6 +344,7 @@ class FusedMoE(CustomOp):
344 tp_size: int | None = None,
345 ep_size: int | None = None,
346 dp_size: int | None = None,
347 + pcp_size: int | None = None,
348 prefix: str = "",
349 custom_routing_function: Callable | None = None,
350 scoring_func: str = "softmax",
@@ -398,12 +400,14 @@ class FusedMoE(CustomOp):
400 tp_size if tp_size is not None else get_tensor_model_parallel_world_size()
401 )
402 dp_size_ = dp_size if dp_size is not None else get_dp_group().world_size
403 + pcp_size_ = pcp_size if pcp_size is not None else get_pcp_group().world_size
404
405 self.is_sequence_parallel = is_sequence_parallel
406 self.sp_size = tp_size_ if is_sequence_parallel else 1
407
408 self.moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(
409 tp_size_=tp_size_,
410 + pcp_size_=pcp_size_,
411 dp_size_=dp_size_,
412 vllm_parallel_config=vllm_config.parallel_config,
413 )
@@ -679,6 +683,10 @@ class FusedMoE(CustomOp):
683 def dp_size(self):
684 return self.moe_parallel_config.dp_size
685
686 + @property
687 + def pcp_size(self):
688 + return self.moe_parallel_config.pcp_size
689 +
690 @property
691 def ep_size(self):
692 return self.moe_parallel_config.ep_size
@@ -691,6 +699,10 @@ class FusedMoE(CustomOp):
699 def dp_rank(self):
700 return self.moe_parallel_config.dp_rank
701
702 + @property
703 + def pcp_rank(self):
704 + return self.moe_parallel_config.pcp_rank
705 +
706 @property
707 def ep_rank(self):
708 return self.moe_parallel_config.ep_rank
@@ -1871,6 +1883,19 @@ class FusedMoE(CustomOp):
1883 assert self.shared_experts is not None
1884 shared_output = self.shared_experts(hidden_states)
1885
1886 + # NOTE: Similar with DP, PCP also needs dispatch and combine. For
1887 + # simplicity, AgRsAll2All was added separately for PCP here. Maybe
1888 + # we should modify All2AllManager abstract to better support PCP.
1889 + if self.pcp_size > 1:
1890 + hidden_states = get_pcp_group().all_gather(
1891 + hidden_states,
1892 + dim=0,
1893 + )
1894 + router_logits = get_pcp_group().all_gather(
1895 + router_logits,
1896 + dim=0,
1897 + )
1898 +
1899 # Matrix multiply.
1900 final_hidden_states = self.quant_method.apply(
1901 layer=self,
@@ -1925,6 +1950,13 @@ class FusedMoE(CustomOp):
1950 def combine_output(states: torch.Tensor) -> torch.Tensor:
1951 if do_naive_dispatch_combine:
1952 states = get_ep_group().combine(states, self.is_sequence_parallel)
1953 +
1954 + if self.pcp_size > 1:
1955 + states = get_pcp_group().reduce_scatter(
1956 + states,
1957 + dim=0,
1958 + )
1959 +
1960 return states
1961
1962 if self.shared_experts is not None:
vllm/model_executor/models/gpt_oss.py
+7 -2
@@ -13,6 +13,7 @@ from vllm.config import CacheConfig, VllmConfig
13 from vllm.distributed import (
14 get_dp_group,
15 get_ep_group,
16 + get_pcp_group,
17 get_pp_group,
18 get_tensor_model_parallel_rank,
19 get_tensor_model_parallel_world_size,
@@ -322,10 +323,12 @@ class GptOssModel(nn.Module):
323
324 # In MoE, we need to flatten the tensor parallel size across the data
325 # parallel size when EP is disabled.
325 - tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp(
326 + tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
327 tp_size=get_tensor_model_parallel_world_size(),
328 dp_size=get_dp_group().world_size,
329 dp_rank=get_dp_group().rank_in_group,
330 + pcp_size=get_pcp_group().world_size,
331 + pcp_rank=get_pcp_group().rank_in_group,
332 )
333
334 intermediate_size = self.config.intermediate_size
@@ -507,10 +510,12 @@ class GptOssModel(nn.Module):
510
511 # In MoE, we need to flatten the tensor parallel size across the data
512 # parallel size when EP is disabled.
510 - tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp(
513 + tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
514 tp_size=get_tensor_model_parallel_world_size(),
515 dp_size=get_dp_group().world_size,
516 dp_rank=get_dp_group().rank_in_group,
517 + pcp_size=get_pcp_group().world_size,
518 + pcp_rank=get_pcp_group().rank_in_group,
519 )
520
521 intermediate_size = self.config.intermediate_size
vllm/v1/attention/backends/flash_attn.py
+3 -3
@@ -265,8 +265,8 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
265 self.dcp_world_size = 1
266 self.dcp_rank = 0
267
268 - self.dcp_kv_cache_interleave_size = (
269 - self.parallel_config.dcp_kv_cache_interleave_size
268 + self.cp_kv_cache_interleave_size = (
269 + self.parallel_config.cp_kv_cache_interleave_size
270 )
271
272 self.use_full_cuda_graph = (
@@ -388,7 +388,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
388 dcp_context_kv_lens_cpu,
389 self.dcp_world_size,
390 self.dcp_rank,
391 - self.dcp_kv_cache_interleave_size,
391 + self.cp_kv_cache_interleave_size,
392 )
393 dcp_context_kv_lens = dcp_context_kv_lens_cpu.to(self.device)
394 max_dcp_context_kv_len = dcp_context_kv_lens.max().item()
vllm/v1/attention/backends/mla/common.py
+3 -3
@@ -536,7 +536,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
536 # DCP might not be initialized in testing
537 self.dcp_world_size = 1
538 self.dcp_rank = 0
539 - self.dcp_local_block_size = parallel_config.dcp_kv_cache_interleave_size
539 + self.dcp_local_block_size = parallel_config.cp_kv_cache_interleave_size
540 self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size
541
542 # Don't try to access the runner on AMD
@@ -1289,8 +1289,8 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1289 get_current_vllm_config()
1290 )
1291 )
1292 - self.dcp_kv_cache_interleave_size: int = (
1293 - get_current_vllm_config().parallel_config.dcp_kv_cache_interleave_size
1292 + self.cp_kv_cache_interleave_size: int = (
1293 + get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size
1294 )
1295
1296 def _flash_attn_varlen_diff_headdims(
vllm/v1/attention/backends/utils.py
+9 -9
@@ -1080,9 +1080,9 @@ def compute_causal_conv1d_metadata(query_start_loc_p: torch.Tensor):
1080
1081 def get_dcp_local_seq_lens(
1082 seq_lens: torch.Tensor,
1083 - dcp_world_size: int = 1,
1083 + dcp_size: int = 1,
1084 dcp_rank: int | None = None,
1085 - dcp_kv_cache_interleave_size: int = 1,
1085 + cp_kv_cache_interleave_size: int = 1,
1086 ) -> torch.Tensor:
1087 """While using dcp, kv_cache size stored on each rank may be different,
1088 use this function to calculate split decode seq_lens of each dcp rank.
@@ -1091,7 +1091,7 @@ def get_dcp_local_seq_lens(
1091 num_requests = seq_lens.size(0)
1092 if dcp_rank is None:
1093 rank_offsets = (
1094 - torch.arange(dcp_world_size, dtype=torch.int32)
1094 + torch.arange(dcp_size, dtype=torch.int32)
1095 .unsqueeze(0)
1096 .repeat(num_requests, 1)
1097 )
@@ -1102,15 +1102,15 @@ def get_dcp_local_seq_lens(
1102 )
1103 base = (
1104 seq_lens_tiled
1105 - // dcp_kv_cache_interleave_size
1106 - // dcp_world_size
1107 - * dcp_kv_cache_interleave_size
1105 + // cp_kv_cache_interleave_size
1106 + // dcp_size
1107 + * cp_kv_cache_interleave_size
1108 )
1109 - remainder = seq_lens_tiled - base * dcp_world_size
1109 + remainder = seq_lens_tiled - base * dcp_size
1110 remainder = torch.clip(
1111 - remainder - rank_offsets * dcp_kv_cache_interleave_size,
1111 + remainder - rank_offsets * cp_kv_cache_interleave_size,
1112 0,
1113 - dcp_kv_cache_interleave_size,
1113 + cp_kv_cache_interleave_size,
1114 )
1115 dcp_local_seq_lens = base + remainder
1116 return dcp_local_seq_lens.squeeze(1)
vllm/v1/core/kv_cache_coordinator.py
+17
@@ -27,6 +27,7 @@ class KVCacheCoordinator(ABC):
27 enable_caching: bool,
28 enable_kv_cache_events: bool,
29 dcp_world_size: int,
30 + pcp_world_size: int,
31 ):
32 self.kv_cache_config = kv_cache_config
33 self.max_model_len = max_model_len
@@ -44,6 +45,7 @@ class KVCacheCoordinator(ABC):
45 block_pool=self.block_pool,
46 kv_cache_group_id=i,
47 dcp_world_size=dcp_world_size,
48 + pcp_world_size=pcp_world_size,
49 )
50 for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups)
51 )
@@ -210,6 +212,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
212 use_eagle: bool,
213 enable_kv_cache_events: bool,
214 dcp_world_size: int,
215 + pcp_world_size: int,
216 ):
217 super().__init__(
218 kv_cache_config,
@@ -218,6 +221,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
221 False,
222 enable_kv_cache_events,
223 dcp_world_size=dcp_world_size,
224 + pcp_world_size=pcp_world_size,
225 )
226 self.num_single_type_manager = len(self.single_type_managers)
227
@@ -250,6 +254,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
254 enable_caching: bool,
255 enable_kv_cache_events: bool,
256 dcp_world_size: int,
257 + pcp_world_size: int,
258 ):
259 super().__init__(
260 kv_cache_config,
@@ -258,12 +263,16 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
263 enable_caching,
264 enable_kv_cache_events,
265 dcp_world_size=dcp_world_size,
266 + pcp_world_size=pcp_world_size,
267 )
268 self.kv_cache_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec
269 self.block_size = self.kv_cache_spec.block_size
270 self.dcp_world_size = dcp_world_size
271 + self.pcp_world_size = pcp_world_size
272 if dcp_world_size > 1:
273 self.block_size *= dcp_world_size
274 + if pcp_world_size > 1:
275 + self.block_size *= pcp_world_size
276 assert len(self.kv_cache_config.kv_cache_groups) == 1, (
277 "UnitaryKVCacheCoordinator assumes only one kv cache group"
278 )
@@ -281,6 +290,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
290 kv_cache_spec=self.kv_cache_spec,
291 use_eagle=self.use_eagle,
292 dcp_world_size=self.dcp_world_size,
293 + pcp_world_size=self.pcp_world_size,
294 )
295 return hit_blocks, len(hit_blocks[0]) * self.block_size
296
@@ -302,6 +312,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
312 enable_caching: bool,
313 enable_kv_cache_events: bool,
314 dcp_world_size: int,
315 + pcp_world_size: int,
316 ):
317 super().__init__(
318 kv_cache_config,
@@ -310,8 +321,10 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
321 enable_caching,
322 enable_kv_cache_events,
323 dcp_world_size=dcp_world_size,
324 + pcp_world_size=pcp_world_size,
325 )
326 assert dcp_world_size == 1, "DCP not support hybrid attn now."
327 + assert pcp_world_size == 1, "PCP not support hybrid attn now."
328 self.verify_and_split_kv_cache_groups()
329
330 def verify_and_split_kv_cache_groups(self) -> None:
@@ -452,6 +465,7 @@ def get_kv_cache_coordinator(
465 enable_caching: bool,
466 enable_kv_cache_events: bool,
467 dcp_world_size: int,
468 + pcp_world_size: int,
469 ) -> KVCacheCoordinator:
470 if not enable_caching:
471 return KVCacheCoordinatorNoPrefixCache(
@@ -460,6 +474,7 @@ def get_kv_cache_coordinator(
474 use_eagle,
475 enable_kv_cache_events,
476 dcp_world_size=dcp_world_size,
477 + pcp_world_size=pcp_world_size,
478 )
479 if len(kv_cache_config.kv_cache_groups) == 1:
480 return UnitaryKVCacheCoordinator(
@@ -469,6 +484,7 @@ def get_kv_cache_coordinator(
484 enable_caching,
485 enable_kv_cache_events,
486 dcp_world_size=dcp_world_size,
487 + pcp_world_size=pcp_world_size,
488 )
489 return HybridKVCacheCoordinator(
490 kv_cache_config,
@@ -477,4 +493,5 @@ def get_kv_cache_coordinator(
493 enable_caching,
494 enable_kv_cache_events,
495 dcp_world_size=dcp_world_size,
496 + pcp_world_size=pcp_world_size,
497 )
vllm/v1/core/kv_cache_manager.py
+4 -5
@@ -100,6 +100,7 @@ class KVCacheManager:
100 log_stats: bool = False,
101 enable_kv_cache_events: bool = False,
102 dcp_world_size: int = 1,
103 + pcp_world_size: int = 1,
104 ) -> None:
105 self.max_model_len = max_model_len
106
@@ -124,12 +125,9 @@ class KVCacheManager:
125 0
126 ].kv_cache_spec.block_size
127
127 - if dcp_world_size > 1:
128 + if dcp_world_size * pcp_world_size > 1:
129 assert len(kv_cache_config.kv_cache_groups) == 1
129 - # Note(hc): need revisit. When both DCP and any future
130 - # PCP are enabled, the block_size may need to be scaled
131 - # by a factor of dcp_size × pcp_size?
132 - self.block_size *= dcp_world_size
130 + self.block_size *= dcp_world_size * pcp_world_size
131
132 self.coordinator = get_kv_cache_coordinator(
133 kv_cache_config=kv_cache_config,
@@ -138,6 +136,7 @@ class KVCacheManager:
136 enable_caching=self.enable_caching,
137 enable_kv_cache_events=enable_kv_cache_events,
138 dcp_world_size=dcp_world_size,
139 + pcp_world_size=pcp_world_size,
140 )
141 self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
142 self.block_pool = self.coordinator.block_pool
vllm/v1/core/kv_cache_utils.py
+9 -4
@@ -1219,11 +1219,16 @@ def _report_kv_cache_config(
1219 // len(kv_cache_config.kv_cache_groups)
1220 * min_block_size
1221 )
1222 - if vllm_config.parallel_config.decode_context_parallel_size > 1:
1223 - num_tokens *= vllm_config.parallel_config.decode_context_parallel_size
1222 + dcp_size = vllm_config.parallel_config.decode_context_parallel_size
1223 + pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
1224 + if pcp_size * dcp_size > 1:
1225 + num_tokens *= pcp_size * dcp_size
1226 logger.info(
1225 - "Multiplying the GPU KV cache size by the dcp_world_size %d.",
1226 - vllm_config.parallel_config.decode_context_parallel_size,
1227 + "Multiplying the GPU KV cache size by the cp_world_size %d "
1228 + "(pcp_world_size %d * dcp_world_size %d).",
1229 + pcp_size * dcp_size,
1230 + pcp_size,
1231 + dcp_size,
1232 )
1233 num_tokens_str = f"{num_tokens:,}"
1234 logger.info_once("GPU KV cache size: %s tokens", num_tokens_str, scope="local")
vllm/v1/core/sched/scheduler.py
+2
@@ -121,6 +121,7 @@ class Scheduler(SchedulerInterface):
121
122 self.block_size = block_size
123 self.dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
124 + self.pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
125
126 # req_id -> Request
127 self.requests: dict[str, Request] = {}
@@ -183,6 +184,7 @@ class Scheduler(SchedulerInterface):
184 log_stats=self.log_stats,
185 enable_kv_cache_events=self.enable_kv_cache_events,
186 dcp_world_size=self.dcp_world_size,
187 + pcp_world_size=self.pcp_world_size,
188 )
189 self.use_pp = self.parallel_config.pipeline_parallel_size > 1
190
vllm/v1/core/single_type_kv_cache_manager.py
+15 -4
@@ -32,6 +32,7 @@ class SingleTypeKVCacheManager(ABC):
32 block_pool: BlockPool,
33 kv_cache_group_id: int,
34 dcp_world_size: int = 1,
35 + pcp_world_size: int = 1,
36 ) -> None:
37 """
38 Initializes the SingleTypeKVCacheManager.
@@ -42,8 +43,9 @@ class SingleTypeKVCacheManager(ABC):
43 """
44 self.block_size = kv_cache_spec.block_size
45 self.dcp_world_size = dcp_world_size
45 - if self.dcp_world_size > 1:
46 - self.block_size *= dcp_world_size
46 + self.pcp_world_size = pcp_world_size
47 + if dcp_world_size * pcp_world_size > 1:
48 + self.block_size *= dcp_world_size * pcp_world_size
49 self.kv_cache_spec = kv_cache_spec
50 self.block_pool = block_pool
51
@@ -212,6 +214,7 @@ class SingleTypeKVCacheManager(ABC):
214 kv_cache_spec: KVCacheSpec,
215 use_eagle: bool,
216 dcp_world_size: int = 1,
217 + pcp_world_size: int = 1,
218 ) -> tuple[list[KVCacheBlock], ...]:
219 """
220 Get the longest cache hit prefix of the blocks that is not longer than
@@ -303,6 +306,7 @@ class FullAttentionManager(SingleTypeKVCacheManager):
306 kv_cache_spec: KVCacheSpec,
307 use_eagle: bool,
308 dcp_world_size: int = 1,
309 + pcp_world_size: int = 1,
310 ) -> tuple[list[KVCacheBlock], ...]:
311 assert isinstance(
312 kv_cache_spec, (FullAttentionSpec, ChunkedLocalAttentionSpec)
@@ -314,8 +318,8 @@ class FullAttentionManager(SingleTypeKVCacheManager):
318 [] for _ in range(len(kv_cache_group_ids))
319 )
320 block_size = kv_cache_spec.block_size
317 - if dcp_world_size > 1:
318 - block_size *= dcp_world_size
321 + if dcp_world_size * pcp_world_size > 1:
322 + block_size *= dcp_world_size * pcp_world_size
323 max_num_blocks = max_length // block_size
324 for block_hash in itertools.islice(block_hashes, max_num_blocks):
325 # block_hashes is a chain of block hashes. If a block hash is not
@@ -362,11 +366,13 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
366 kv_cache_spec: KVCacheSpec,
367 use_eagle: bool,
368 dcp_world_size: int = 1,
369 + pcp_world_size: int = 1,
370 ) -> tuple[list[KVCacheBlock], ...]:
371 assert isinstance(kv_cache_spec, SlidingWindowSpec), (
372 "SlidingWindowManager can only be used for sliding window groups"
373 )
374 assert dcp_world_size == 1, "DCP not support sliding window attn now."
375 + assert pcp_world_size == 1, "PCP not support sliding window attn now."
376
377 # The number of contiguous blocks needed for prefix cache hit.
378 # -1 since the input token itself is also included in the window
@@ -476,6 +482,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
482 kv_cache_spec: KVCacheSpec,
483 use_eagle: bool,
484 dcp_world_size: int = 1,
485 + pcp_world_size: int = 1,
486 ) -> tuple[list[KVCacheBlock], ...]:
487 """
488 For chunked local attention, we need to find the longest cache hit
@@ -516,6 +523,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
523 "Hybrid KV cache is not supported for " + "eagle + chunked local attention."
524 )
525 assert dcp_world_size == 1, "DCP not support chunked local attn now."
526 + assert pcp_world_size == 1, "PCP not support chunked local attn now."
527 max_num_blocks = max_length // kv_cache_spec.block_size
528 if max_length > 0:
529 local_attention_start_idx = (
@@ -611,11 +619,13 @@ class MambaManager(SingleTypeKVCacheManager):
619 kv_cache_spec: KVCacheSpec,
620 use_eagle: bool,
621 dcp_world_size: int = 1,
622 + pcp_world_size: int = 1,
623 ) -> tuple[list[KVCacheBlock], ...]:
624 assert isinstance(kv_cache_spec, MambaSpec), (
625 "MambaManager can only be used for mamba groups"
626 )
627 assert dcp_world_size == 1, "DCP not support mamba now."
628 + assert pcp_world_size == 1, "PCP not support mamba now."
629 computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
630 [] for _ in range(len(kv_cache_group_ids))
631 )
@@ -705,6 +715,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
715 kv_cache_spec: KVCacheSpec,
716 use_eagle: bool,
717 dcp_world_size: int = 1,
718 + pcp_world_size: int = 1,
719 ) -> tuple[list[KVCacheBlock], ...]:
720 assert isinstance(kv_cache_spec, CrossAttentionSpec), (
721 "CrossAttentionManager can only be used for cross-attention groups"
vllm/v1/engine/core.py
+1
@@ -128,6 +128,7 @@ class EngineCore:
128 scheduler_block_size = (
129 vllm_config.cache_config.block_size
130 * vllm_config.parallel_config.decode_context_parallel_size
131 + * vllm_config.parallel_config.prefill_context_parallel_size
132 )
133
134 self.scheduler: SchedulerInterface = Scheduler(
vllm/v1/executor/multiproc_executor.py
+17 -6
@@ -35,6 +35,7 @@ from vllm.distributed.parallel_state import (
35 get_dp_group,
36 get_ep_group,
37 get_inner_dp_world_group,
38 + get_pcp_group,
39 get_pp_group,
40 get_tp_group,
41 )
@@ -110,12 +111,14 @@ class MultiprocExecutor(Executor):
111 f"({self.parallel_config.nnodes_within_dp}). "
112 )
113 self.local_world_size = self.parallel_config.local_world_size
113 - tensor_parallel_size = self.parallel_config.tensor_parallel_size
114 - pp_parallel_size = self.parallel_config.pipeline_parallel_size
115 - assert self.world_size == tensor_parallel_size * pp_parallel_size, (
114 + tp_size = self.parallel_config.tensor_parallel_size
115 + pp_size = self.parallel_config.pipeline_parallel_size
116 + pcp_size = self.parallel_config.prefill_context_parallel_size
117 + assert self.world_size == tp_size * pp_size * pcp_size, (
118 f"world_size ({self.world_size}) must be equal to the "
117 - f"tensor_parallel_size ({tensor_parallel_size}) x pipeline"
118 - f"_parallel_size ({pp_parallel_size}). "
119 + f"tensor_parallel_size ({tp_size}) x pipeline"
120 + f"_parallel_size ({pp_size}) x prefill_context"
121 + f"_parallel_size ({pcp_size}). "
122 )
123
124 # Set multiprocessing envs
@@ -424,7 +427,11 @@ class MultiprocExecutor(Executor):
427 # 16-23, PP rank 2
428 # 24-31, PP rank 3
429 # so world_size - tp_size = 32 - 8 = 24 should be PP rank = -1 (i.e. 3)
427 - return self.world_size - self.parallel_config.tensor_parallel_size
430 + return (
431 + self.world_size
432 + - self.parallel_config.tensor_parallel_size
433 + * self.parallel_config.prefill_context_parallel_size
434 + )
435
436
437 @dataclass
@@ -828,6 +835,8 @@ class WorkerProc:
835 dp_rank = get_dp_group().rank_in_group
836 pp_size = get_pp_group().world_size
837 pp_rank = get_pp_group().rank_in_group
838 + pcp_size = get_pcp_group().world_size
839 + pcp_rank = get_pcp_group().rank_in_group
840 tp_size = get_tp_group().world_size
841 tp_rank = get_tp_group().rank_in_group
842 dcp_size = get_dcp_group().world_size
@@ -837,6 +846,8 @@ class WorkerProc:
846 process_name += f"_DP{dp_rank}"
847 if pp_size > 1:
848 process_name += f"_PP{pp_rank}"
849 + if pcp_size > 1:
850 + process_name += f"_PCP{pcp_rank}"
851 if tp_size > 1:
852 process_name += f"_TP{tp_rank}"
853 if dcp_size > 1:
vllm/v1/kv_cache_interface.py
+3 -2
@@ -95,10 +95,11 @@ class FullAttentionSpec(AttentionSpec):
95 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
96 max_model_len = vllm_config.model_config.max_model_len
97 dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
98 + pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
99 # Note(hc): each dcp rank only need save
100 # (max_model_len//dcp_world_size) tokens locally.
100 - if dcp_world_size > 1:
101 - max_model_len = cdiv(max_model_len, dcp_world_size)
101 + if dcp_world_size * pcp_world_size > 1:
102 + max_model_len = cdiv(max_model_len, dcp_world_size * pcp_world_size)
103 return cdiv(max_model_len, self.block_size) * self.page_size_bytes
104
105 @classmethod
vllm/v1/worker/block_table.py
+22 -13
@@ -4,7 +4,7 @@
4 import numpy as np
5 import torch
6
7 -from vllm.distributed import get_dcp_group
7 +from vllm.distributed import get_dcp_group, get_pcp_group
8 from vllm.logger import init_logger
9 from vllm.utils.math_utils import cdiv
10 from vllm.v1.utils import CpuGpuBuffer
@@ -22,7 +22,7 @@ class BlockTable:
22 pin_memory: bool,
23 device: torch.device,
24 kernel_block_size: int,
25 - dcp_kv_cache_interleave_size: int,
25 + cp_kv_cache_interleave_size: int,
26 ):
27 """
28 Args:
@@ -80,6 +80,13 @@ class BlockTable:
80 else:
81 self._kernel_block_arange = None
82
83 + try:
84 + self.pcp_world_size = get_pcp_group().world_size
85 + self.pcp_rank = get_pcp_group().rank_in_group
86 + except AssertionError:
87 + # DCP might not be initialized in testing
88 + self.pcp_world_size = 1
89 + self.pcp_rank = 0
90 try:
91 self.dcp_world_size = get_dcp_group().world_size
92 self.dcp_rank = get_dcp_group().rank_in_group
@@ -87,7 +94,7 @@ class BlockTable:
94 # DCP might not be initialized in testing
95 self.dcp_world_size = 1
96 self.dcp_rank = 0
90 - self.dcp_kv_cache_interleave_size = dcp_kv_cache_interleave_size
97 + self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size
98
99 def append_row(
100 self,
@@ -131,14 +138,16 @@ class BlockTable:
138 # NOTE(woosuk): We can't simply use `token_indices // block_size`
139 # here because M (max_model_len) is not necessarily divisible by
140 # block_size.
134 - if self.dcp_world_size > 1:
141 + total_cp_world_size = self.pcp_world_size * self.dcp_world_size
142 + total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank
143 + if total_cp_world_size > 1:
144 # Note(hc): The DCP implement store kvcache with an interleave
145 # style, the kvcache for the token whose token_idx is i is
146 # always stored on the GPU whose dcp_rank equals i % cp_world_size:
147
148 # Use a "virtual block" which equals to world_size * block_size
149 # for block_table_indices calculation.
141 - virtual_block_size = self.block_size * self.dcp_world_size
150 + virtual_block_size = self.block_size * total_cp_world_size
151 block_table_indices = (
152 req_indices * self.max_num_blocks_per_req
153 + positions // virtual_block_size
@@ -150,16 +159,16 @@ class BlockTable:
159 virtual_block_offsets = positions % virtual_block_size
160 mask = (
161 virtual_block_offsets
153 - // self.dcp_kv_cache_interleave_size
154 - % self.dcp_world_size
155 - == self.dcp_rank
162 + // self.cp_kv_cache_interleave_size
163 + % total_cp_world_size
164 + == total_cp_rank
165 )
166 # Calculate local block_offsets
167 block_offsets = (
168 virtual_block_offsets
160 - // (self.dcp_world_size * self.dcp_kv_cache_interleave_size)
161 - * self.dcp_kv_cache_interleave_size
162 - + virtual_block_offsets % self.dcp_kv_cache_interleave_size
169 + // (total_cp_world_size * self.cp_kv_cache_interleave_size)
170 + * self.cp_kv_cache_interleave_size
171 + + virtual_block_offsets % self.cp_kv_cache_interleave_size
172 )
173 # Calculate slot_mapping
174 slot_mapping = block_numbers * self.block_size + block_offsets
@@ -253,7 +262,7 @@ class MultiGroupBlockTable:
262 block_sizes: list[int],
263 kernel_block_sizes: list[int],
264 num_speculative_tokens: int = 0,
256 - dcp_kv_cache_interleave_size: int = 1,
265 + cp_kv_cache_interleave_size: int = 1,
266 ) -> None:
267 # Note(hc): each dcp rank only store
268 # (max_model_len//dcp_world_size) tokens in kvcache,
@@ -283,7 +292,7 @@ class MultiGroupBlockTable:
292 pin_memory,
293 device,
294 kernel_block_size,
286 - dcp_kv_cache_interleave_size,
295 + cp_kv_cache_interleave_size,
296 )
297 for block_size, kernel_block_size in zip(block_sizes, kernel_block_sizes)
298 ]
vllm/v1/worker/gpu_input_batch.py
+2 -2
@@ -87,7 +87,7 @@ class InputBatch:
87 is_spec_decode: bool = False,
88 is_pooling_model: bool = False,
89 num_speculative_tokens: int = 0,
90 - dcp_kv_cache_interleave_size: int = 1,
90 + cp_kv_cache_interleave_size: int = 1,
91 ):
92 self.is_pooling_model = is_pooling_model
93 self.is_spec_decode = is_spec_decode
@@ -141,7 +141,7 @@ class InputBatch:
141 block_sizes=block_sizes,
142 kernel_block_sizes=kernel_block_sizes,
143 num_speculative_tokens=num_speculative_tokens,
144 - dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,
144 + cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
145 )
146
147 # Sampling-related.
vllm/v1/worker/gpu_model_runner.py
+2 -2
@@ -426,7 +426,7 @@ class GPUModelRunner(
426 # uses output token ids so we set this conservatively.
427 logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
428 is_pooling_model=self.is_pooling_model,
429 - dcp_kv_cache_interleave_size=self.parallel_config.dcp_kv_cache_interleave_size,
429 + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
430 )
431
432 self.use_async_scheduling = self.scheduler_config.async_scheduling
@@ -1436,7 +1436,7 @@ class GPUModelRunner(
1436 self.seq_lens.cpu[:num_reqs],
1437 self.dcp_world_size,
1438 self.dcp_rank,
1439 - self.parallel_config.dcp_kv_cache_interleave_size,
1439 + self.parallel_config.cp_kv_cache_interleave_size,
1440 )
1441 self.dcp_local_seq_lens.copy_to_gpu(num_reqs)
1442
vllm/v1/worker/gpu_worker.py
+3
@@ -26,6 +26,7 @@ from vllm.distributed.kv_transfer import (
26 has_kv_transfer_group,
27 )
28 from vllm.distributed.parallel_state import (
29 + get_pcp_group,
30 get_pp_group,
31 get_tp_group,
32 )
@@ -733,6 +734,7 @@ class Worker(WorkerBase):
734 module.global_num_experts = module.moe_config.num_experts
735 module.moe_parallel_config = FusedMoEParallelConfig.make(
736 tp_size_=get_tp_group().world_size,
737 + pcp_size_=get_pcp_group().world_size,
738 dp_size_=get_dp_group().world_size,
739 vllm_parallel_config=parallel_config,
740 )
@@ -886,6 +888,7 @@ def init_worker_distributed_environment(
888 ensure_model_parallel_initialized(
889 parallel_config.tensor_parallel_size,
890 parallel_config.pipeline_parallel_size,
891 + parallel_config.prefill_context_parallel_size,
892 parallel_config.decode_context_parallel_size,
893 )
894