@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b6ff8a2f50

[Core] Add MRV2 virtual-batch PCP for MLA (#46570)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

Lucas Wilkinson committed Jul 19, 2026 at 04:53 UTC b6ff8a2f509cc7ac9c58176f5115a836aa1e08bd

39 files changed +1204 -196

.buildkite/test_areas/lm_eval.yaml

+22

@@ -78,6 +78,28 @@ steps:

78 commands:

79 - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt

81 +- label: LM Eval PCP (4xB200)

82 + key: lm-eval-pcp-4xb200

83 + timeout_in_minutes: 360

84 + device: b200-k8s

85 + num_devices: 4

86 + optional: true

87 + source_file_dependencies:

88 + - csrc/

89 + - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml

90 + - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml

91 + - tests/evals/gsm8k/configs/models-pcp.txt

92 + - vllm/model_executor/layers/quantization

93 + - vllm/config/parallel.py

94 + - vllm/distributed/parallel_state.py

95 + - vllm/model_executor/layers/attention/mla_attention.py

96 + - vllm/model_executor/layers/attention/pcp.py

97 + - vllm/v1/worker/gpu/model_runner.py

98 + - vllm/v1/worker/gpu/pcp_manager.py

99 + autorun_on_main: true

100 + commands:

101 + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt

102 +

103 - label: LM Eval Large Models EP (2xB200)

104 key: lm-eval-large-models-ep-2xb200

105 timeout_in_minutes: 60

tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml new

+16

@@ -0,0 +1,16 @@

1 +model_name: "nvidia/GLM-5.2-NVFP4"

2 +accuracy_threshold: 0.90

3 +num_questions: 1319

4 +num_fewshot: 5

5 +max_concurrency: 100

6 +server_args: >-

7 + --enforce-eager

8 + --max-model-len 4096

9 + --safetensors-load-strategy prefetch

10 + --moe-backend flashinfer_cutlass

11 + --prefill-context-parallel-size 4

12 + --enable-expert-parallel

13 + --kv-cache-dtype fp8

14 +env:

15 + VLLM_LOGGING_LEVEL: "DEBUG"

16 + VLLM_USE_V2_MODEL_RUNNER: "1"

tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml new

+17

@@ -0,0 +1,17 @@

1 +model_name: "nvidia/GLM-5.2-NVFP4"

2 +accuracy_threshold: 0.90

3 +num_questions: 1319

4 +num_fewshot: 5

5 +max_concurrency: 100

6 +server_args: >-

7 + --enforce-eager

8 + --max-model-len 4096

9 + --safetensors-load-strategy prefetch

10 + --moe-backend flashinfer_cutlass

11 + --tensor-parallel-size 2

12 + --prefill-context-parallel-size 2

13 + --enable-expert-parallel

14 + --kv-cache-dtype fp8

15 +env:

16 + VLLM_LOGGING_LEVEL: "DEBUG"

17 + VLLM_USE_V2_MODEL_RUNNER: "1"

tests/evals/gsm8k/configs/models-pcp.txt new

@@ -0,0 +1,2 @@

1 +GLM-5.2-NVFP4-TP2-PCP2-EP.yaml

2 +GLM-5.2-NVFP4-TP1-PCP4-EP.yaml

tests/evals/gsm8k/gsm8k_eval.py

+15 -1

@@ -217,6 +217,7 @@ def evaluate_gsm8k(

217 seed: int | None = 42,

218 request_timeout_seconds: float = 600,

219 gen_prefix: str = "",

220 + max_concurrency: int | None = None,

221 ) -> dict[str, float | int]:

222 """

223 Evaluate GSM8K accuracy using vLLM serve endpoint.

@@ -261,7 +262,14 @@ def evaluate_gsm8k(

262 return answer, tokens

263

264 timeout = aiohttp.ClientTimeout(total=request_timeout_seconds)

264 - async with aiohttp.ClientSession(timeout=timeout) as session:

265 + connector = (

266 + aiohttp.TCPConnector(limit=max_concurrency)

267 + if max_concurrency is not None

268 + else None

269 + )

270 + async with aiohttp.ClientSession(

271 + timeout=timeout, connector=connector

272 + ) as session:

273 tasks = [get_answer(session, i) for i in range(num_questions)]

274 await tqdm.gather(*tasks, desc="Evaluating")

275

@@ -343,6 +351,11 @@ def main() -> None:

351 parser.add_argument(

352 "--seed", type=int, default=42, help="Random seed for reproducibility"

353 )

354 + parser.add_argument(

355 + "--max-concurrency",

356 + type=int,

357 + help="Maximum number of concurrent requests",

358 + )

359 parser.add_argument("--save-results", type=str, help="Save results to JSON file")

360

361 args = parser.parse_args()

@@ -355,6 +368,7 @@ def main() -> None:

368 port=args.port,

369 temperature=args.temperature,

370 seed=args.seed,

371 + max_concurrency=args.max_concurrency,

372 )

373

374 # Print results to terminal

tests/evals/gsm8k/test_gsm8k_correctness.py

@@ -59,6 +59,7 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict:

59 seed=eval_config.get("seed", 42),

60 request_timeout_seconds=request_timeout_seconds,

61 gen_prefix=eval_config.get("gen_prefix", ""),

62 + max_concurrency=eval_config.get("max_concurrency"),

63 )

65 return results

tests/v1/attention/test_attention_splitting.py

+16

@@ -160,6 +160,8 @@ def apply_split_decodes_and_prefills(

160 decode_threshold: int,

161 require_uniform: bool,

162 padded_num_tokens: int | None = None,

163 + is_prefilling: list[bool] | None = None,

164 + treat_short_extends_as_decodes: bool = True,

165 ):

166 """Helper function to apply split_decodes_and_prefills and return

167 the results."""

@@ -173,11 +175,14 @@ def apply_split_decodes_and_prefills(

175

176 if padded_num_tokens is not None:

177 common_metadata.num_actual_tokens = padded_num_tokens

178 + if is_prefilling is not None:

179 + common_metadata.is_prefilling = torch.tensor(is_prefilling)

180

181 return split_decodes_and_prefills(

182 common_metadata,

183 decode_threshold=decode_threshold,

184 require_uniform=require_uniform,

185 + treat_short_extends_as_decodes=treat_short_extends_as_decodes,

186 )

187

188

@@ -236,6 +241,17 @@ def test_split_decodes_and_prefills_uniform_all_ones():

241 assert num_prefill_tokens == 0

242

243

244 +def test_split_decodes_and_prefills_uniform_short_extend():

245 + result = apply_split_decodes_and_prefills(

246 + [1, 1],

247 + decode_threshold=1,

248 + require_uniform=True,

249 + is_prefilling=[False, True],

250 + treat_short_extends_as_decodes=False,

251 + )

252 + assert result == (1, 1, 1, 1)

253 +

254 +

255 def test_split_decodes_and_prefills_uniform_all_short_decodes():

256 query_lens = [2, 2, 1, 3, 2, 1, 2]

257 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

tests/v1/attention/test_sparse_mla_backends.py

@@ -1114,6 +1114,13 @@ def test_sparse_backend_prefill_correctness(

1114 @pytest.mark.parametrize(

1115 "seq_lens,query_lens,workspace_size,max_logits_bytes,expected",

1116 [

1117 + (

1118 + torch.tensor([0]),

1119 + torch.tensor([0]),

1120 + 100,

1121 + 1000,

1122 + [],

1123 + ),

1124 # Logits constraint triggers split (M*N exceeds budget)

1125 # req0: M=10, N=100 -> 1000 elems (4000 bytes) - fits in 5000

1126 # req1: adding M=10, N=100 -> new_M=20, new_N=200 -> 4000 elems > 1250

tests/v1/kv_offload/test_factory.py

+4 -4

@@ -385,7 +385,7 @@ def test_tiering_spec_aligns_row_size():

385 assert spec.num_blocks == cpu_bytes_to_use // alignment

386

387

388 -def test_offloading_spec_resolves_prefill_context_parallel_block_sizes():

388 +def test_offloading_spec_kv_sharding_ignores_prefill_context_parallel():

389 config = _make_layout_vllm_config(

390 cpu_bytes_to_use=65536,

391 extra_config={"block_size": 64},

@@ -394,9 +394,9 @@ def test_offloading_spec_resolves_prefill_context_parallel_block_sizes():

394

395 spec = _create_spec(config, _make_kv_cache_config())

396

397 - assert spec.tokens_per_block == (32,)

398 - assert spec.tokens_per_hash == 32

399 - assert spec.blocks_per_chunk == 2

397 + assert spec.tokens_per_block == (16,)

398 + assert spec.tokens_per_hash == 16

399 + assert spec.blocks_per_chunk == 4

400

401

402 def test_offloading_config_preserves_data_parallel_index():

tests/v1/simple_kv_offload/test_scheduler.py

+16 -24

@@ -1551,9 +1551,8 @@ def _make_cp_scheduler(

1551 num_gpu_blocks: int = 16,

1552 lazy: bool = False,

1553 ) -> SchedulerFixture:

1554 - """Build a SimpleCPUOffloadScheduler with CP-scaled virtual block size."""

1555 - cp_world_size = dcp_world_size * pcp_world_size

1556 - virtual_block_size = BLOCK_SIZE * cp_world_size

1554 + """Build a SimpleCPUOffloadScheduler with DCP-scaled block size."""

1555 + virtual_block_size = BLOCK_SIZE * dcp_world_size

1556

1557 kv_cache_config = _make_kv_cache_config(num_gpu_blocks)

1558 vllm_config = _make_cp_vllm_config(dcp_world_size, pcp_world_size)

@@ -1645,17 +1644,14 @@ def _allocate_cp_gpu_blocks(

1644 ],

1645 )

1646 def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None:

1648 - """Verify that the scheduler's block_size and cp_world_size are correctly

1649 - scaled when context parallelism is enabled."""

1647 + """Verify block size scaling follows DCP ownership."""

1648 fix = _make_cp_scheduler(

1649 dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size

1650 )

1651 sched = fix.scheduler

1652

1655 - expected_cp = dcp_world_size * pcp_world_size

1656 - assert sched.cp_world_size == expected_cp

1657 - assert sched.block_size == BLOCK_SIZE * expected_cp

1658 - assert sched.fa_block_size == BLOCK_SIZE * expected_cp

1653 + assert sched.cp_world_size == dcp_world_size

1654 + assert sched.block_size == BLOCK_SIZE * dcp_world_size

1655

1656

1657 # ---------------------------------------------------------------------------

@@ -1682,8 +1678,7 @@ def test_cp_eager_store_and_load_roundtrip(

1678 lazy=False,

1679 )

1680 sched = fix.scheduler

1685 - cp = dcp_world_size * pcp_world_size

1686 - vbs = BLOCK_SIZE * cp

1681 + vbs = BLOCK_SIZE * dcp_world_size

1682

1683 num_blocks = 2

1684 req = _make_cp_request(num_blocks, vbs)

@@ -1749,18 +1744,16 @@ def test_cp_eager_store_and_load_roundtrip(

1744 def test_cp_effective_block_size_store_and_load(

1745 dcp_world_size: int, pcp_world_size: int

1746 ) -> None:

1752 - """Verify ready_blocks_g (store) and n_take_g (load) use the effective

1753 - (physical * cp) block size, not the per-rank physical size."""

1747 + """Verify store/load physical block counts scale with DCP, not PCP."""

1748 fix = _make_cp_scheduler(

1749 dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size

1750 )

1751 sched = fix.scheduler

1752 gpu_pool = fix.gpu_block_pool

1759 - cp = dcp_world_size * pcp_world_size

1760 - vbs = BLOCK_SIZE * cp

1753 + vbs = BLOCK_SIZE * dcp_world_size

1754 + expected_blocks = 1

1755

1762 - # Store: allocate 2 blocks, confirm only 1. Without the fix,

1763 - # ready_blocks_g = vbs / BLOCK_SIZE = 2, storing both blocks.

1756 + # Store one DCP-scaled logical block worth of tokens.

1757 req = _make_cp_request(num_blocks=2, virtual_block_size=vbs)

1758 gpu_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs)

1759 kv = KVCacheBlocks(blocks=(gpu_blocks,))

@@ -1772,12 +1765,11 @@ def test_cp_effective_block_size_store_and_load(

1765 new_reqs={req.request_id: kv.get_block_ids()},

1766 )

1767 )

1775 - assert len(m1.store_gpu_blocks) == 1

1776 - assert len(m1.store_cpu_blocks) == 1

1768 + assert len(m1.store_gpu_blocks) == expected_blocks

1769 + assert len(m1.store_cpu_blocks) == expected_blocks

1770 simulate_store_completion(sched, m1.store_event)

1771

1779 - # Load: store 2 blocks from a second request, accept only 1 as external.

1780 - # Without the fix, n_take_g = vbs / BLOCK_SIZE = 2, loading both.

1772 + # Load one DCP-scaled logical block from a two-block CPU hit.

1773 req2 = _make_cp_request(num_blocks=2, virtual_block_size=vbs)

1774 kv2 = KVCacheBlocks(blocks=(_allocate_cp_gpu_blocks(gpu_pool, req2, 2, vbs),))

1775 req2.num_computed_tokens = 2 * vbs

@@ -1801,7 +1793,7 @@ def test_cp_effective_block_size_store_and_load(

1793 hit, _ = sched.get_num_new_matched_tokens(req3, num_computed_tokens=0)

1794 assert hit == 2 * vbs

1795

1804 - kv3 = KVCacheBlocks(blocks=(gpu_pool.get_new_blocks(2),))

1796 + kv3 = KVCacheBlocks(blocks=(gpu_pool.get_new_blocks(expected_blocks),))

1797 sched.update_state_after_alloc(req3, kv3, num_external_tokens=vbs)

1798 m3 = sched.build_connector_meta(

1799 make_scheduler_output(

@@ -1810,8 +1802,8 @@ def test_cp_effective_block_size_store_and_load(

1802 )

1803 )

1804 assert m3.load_event >= 0

1813 - assert len(m3.load_gpu_blocks) == 1

1814 - assert len(m3.load_cpu_blocks) == 1

1805 + assert len(m3.load_gpu_blocks) == expected_blocks

1806 + assert len(m3.load_cpu_blocks) == expected_blocks

1807 assert m3.load_gpu_blocks == [kv3.get_block_ids()[0][0]]

1808

1809

vllm/config/parallel.py

+37 -25

@@ -122,10 +122,11 @@ class ParallelConfig:

122 tensor_parallel_size: int = Field(default=1, ge=1)

123 """Number of tensor parallel groups."""

124 prefill_context_parallel_size: int = Field(default=1, ge=1)

125 - """Number of prefill context parallel groups."""

125 + """Number of ranks that split prefill sequence computation. PCP expands

126 + the process world size but does not increase the KV-cache shard count."""

127 data_parallel_size: int = Field(default=1, ge=1)

128 """Number of data parallel groups. MoE layers will be sharded according to

128 - the product of the tensor parallel size and data parallel size."""

129 + the product of the tensor, prefill-context, and data parallel sizes."""

130 data_parallel_size_local: int = Field(default=1, ge=0)

131 """Number of local data parallel groups. A value of 0 is a sentinel used by

132 the engine-args layer to signal that data parallelism was specified

@@ -337,9 +338,9 @@ class ParallelConfig:

338 connect as clients to exchange self-picked group ports at runtime."""

339

340 decode_context_parallel_size: int = Field(default=1, ge=1)

340 - """Number of decode context parallel groups, because the world size does

341 - not change by dcp, it simply reuse the GPUs of TP group, and tp_size

342 - needs to be divisible by dcp_size."""

341 + """Number of ranks that shard the decode KV cache. DCP does not expand

342 + the process world size. Without PCP, DCP reuses TP ranks. With PCP, DCP

343 + either spans the PCP axis or the full TP x PCP block."""

344

345 dcp_kv_cache_interleave_size: int = 1

346 """

@@ -357,13 +358,11 @@ class ParallelConfig:

358 """

359

360 cp_kv_cache_interleave_size: int = 1

360 - """Interleave size of kv_cache storage while using DCP or PCP.

361 - For `total_cp_rank = pcp_rank * dcp_world_size + dcp_rank`,

362 - and `total_cp_world_size = pcp_world_size * dcp_world_size`.

363 - store interleave_size tokens on total_cp_rank i,

364 - then store next interleave_size tokens on total_cp_rank i+1.

361 + """Interleave size of kv_cache storage while using DCP.

362 + Store interleave_size tokens on dcp_rank i, then store next

363 + interleave_size tokens on dcp_rank i+1.

364 Interleave_size=1: token-level alignment, where token `i` is stored on

366 - total_cp_rank `i % total_cp_world_size`.

365 + dcp_rank `i % dcp_world_size`.

366 Interleave_size=block_size: block-level alignment, where tokens are

367 first populated to the preceding ranks. Tokens are then stored

368 in (rank i+1, block j) only after (rank i, block j) is fully occupied.

@@ -480,11 +479,19 @@ class ParallelConfig:

479 )

480 if not self.enable_expert_parallel:

481 raise ValueError("enable_expert_parallel must be True to use EPLB.")

483 - if self.tensor_parallel_size * self.data_parallel_size <= 1:

482 + # The EP group spans the TP x PCP x DP ranks. EPLB therefore needs

483 + # TP, PCP, or DP > 1.

484 + if (

485 + self.tensor_parallel_size

486 + * self.prefill_context_parallel_size

487 + * self.data_parallel_size

488 + <= 1

489 + ):

490 raise ValueError(

485 - "EPLB requires tensor_parallel_size or data_parallel_size "

486 - f"to be greater than 1, but got "

487 - f"TP={self.tensor_parallel_size},DP={self.data_parallel_size}."

491 + "EPLB requires tensor, prefill-context, or data parallelism, "

492 + f"but got TP={self.tensor_parallel_size}, "

493 + f"PCP={self.prefill_context_parallel_size}, "

494 + f"DP={self.data_parallel_size}."

495 )

496 else:

497 if self.eplb_config.num_redundant_experts != 0:

@@ -495,15 +502,21 @@ class ParallelConfig:

502 "num_redundant_experts."

503 )

504

498 - # Note(hc): In the current implementation of decode context

499 - # parallel(DCP), tp_size needs to be divisible by dcp_size,

500 - # because the world size does not change by dcp, it simply

501 - # reuses the GPUs of TP group, and split one TP group into

502 - # tp_size//dcp_size DCP groups.

503 - if self.tensor_parallel_size % self.decode_context_parallel_size != 0:

505 + tp = self.tensor_parallel_size

506 + pcp = self.prefill_context_parallel_size

507 + dcp = self.decode_context_parallel_size

508 + if pcp > 1 and self.data_parallel_size > 1:

509 + raise ValueError("PCP does not support data parallelism yet.")

510 + if pcp == 1:

511 + # DCP reuses the TP ranks when PCP is disabled.

512 + if tp % dcp != 0:

513 + raise ValueError(f"tp_size={tp} must be divisible by dcp_size={dcp}.")

514 + elif dcp not in (1, pcp, tp * pcp):

515 raise ValueError(

505 - f"tp_size={self.tensor_parallel_size} must be divisible by"

506 - f"dcp_size={self.decode_context_parallel_size}."

516 + "When PCP is enabled, DCP must be disabled, span the PCP "

517 + "axis, or span the full TP x PCP axis. "

518 + f"Got TP={tp}, PCP={pcp}, DCP={dcp}; valid DCP sizes are "

519 + f"{sorted({1, pcp, tp * pcp})}."

520 )

521

522 if self.dcp_comm_backend == "a2a" and self.decode_context_parallel_size <= 1:

@@ -515,8 +528,7 @@ class ParallelConfig:

528

529 @property

530 def world_size_across_dp(self) -> int:

518 - """world_size_across_dp is TPxPPxDP, it is the size of the world

519 - including data parallelism."""

531 + """Process world size across TP, PCP, PP, and DP."""

532 return self.world_size * self.data_parallel_size

533

534 @property

vllm/config/vllm.py

+3 -2

@@ -2130,9 +2130,10 @@ class VllmConfig:

2130 model_config = self.model_config

2131 speculative_config = self.speculative_config

2132

2133 - if self.parallel_config.prefill_context_parallel_size > 1:

2133 + if self.parallel_config.prefill_context_parallel_size > 1 and not (

2134 + model_config is not None and model_config.use_mla

2135 + ):

2136 unsupported.append("prefill context parallelism")

2135 -

2137 if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE:

2138 unsupported.append("stock torch.compile")

2139

vllm/distributed/device_communicators/all2all.py

+27 -17

@@ -8,7 +8,7 @@ import torch

8 import torch.distributed as dist

10 import vllm.envs as envs

11 -from vllm.distributed import get_dp_group, get_ep_group

11 +from vllm.distributed import get_dp_group, get_ep_group, get_pcp_group

12 from vllm.distributed.utils import StatelessProcessGroup

13 from vllm.forward_context import get_forward_context

14 from vllm.logger import init_logger

@@ -49,6 +49,23 @@ class AgRsAll2AllManager(All2AllManagerBase):

49 def __init__(self, cpu_group, tcp_store_group=None):

50 super().__init__(cpu_group, tcp_store_group)

52 + def _get_comm_group(self, is_sequence_parallel: bool) -> Any:

53 + if is_sequence_parallel:

54 + return get_ep_group()

55 + if self.dp_world_size > 1:

56 + return get_dp_group()

57 + return get_pcp_group()

58 +

59 + def _get_sizes(self, num_local_tokens: int, comm_group: Any) -> list[int]:

60 + if self.dp_world_size == 1:

61 + return [num_local_tokens] * comm_group.world_size

62 +

63 + dp_metadata = get_forward_context().dp_metadata

64 + assert dp_metadata is not None

65 + sizes = dp_metadata.get_chunk_sizes_across_dp_rank()

66 + assert sizes is not None

67 + return sizes

68 +

69 def dispatch_router_logits(

70 self,

71 hidden_states: torch.Tensor,

@@ -62,11 +79,8 @@ class AgRsAll2AllManager(All2AllManagerBase):

79 """

80 Gather hidden_states and router_logits from all dp ranks.

81 """

65 - dp_metadata = get_forward_context().dp_metadata

66 - assert dp_metadata is not None

67 - sizes = dp_metadata.get_chunk_sizes_across_dp_rank()

68 - assert sizes is not None

69 - dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()

82 + dist_group = self._get_comm_group(is_sequence_parallel)

83 + sizes = self._get_sizes(hidden_states.shape[0], dist_group)

84 assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]

86 tensors_to_gather = [hidden_states, router_logits]

@@ -97,11 +111,8 @@ class AgRsAll2AllManager(All2AllManagerBase):

111 """

112 Gather hidden_states and router_logits from all dp ranks.

113 """

100 - dp_metadata = get_forward_context().dp_metadata

101 - assert dp_metadata is not None

102 - sizes = dp_metadata.get_chunk_sizes_across_dp_rank()

103 - assert sizes is not None

104 - dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()

114 + dist_group = self._get_comm_group(is_sequence_parallel)

115 + sizes = self._get_sizes(hidden_states.shape[0], dist_group)

116 assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]

117

118 tensors_to_gather = [hidden_states, topk_weights, topk_ids]

@@ -129,12 +140,11 @@ class AgRsAll2AllManager(All2AllManagerBase):

140 """

141 Reduce-scatter hidden_states across all dp ranks.

142 """

132 - dp_metadata = get_forward_context().dp_metadata

133 - assert dp_metadata is not None

134 - sizes = dp_metadata.get_chunk_sizes_across_dp_rank()

135 - assert sizes is not None

136 -

137 - dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()

143 + dist_group = self._get_comm_group(is_sequence_parallel)

144 + sizes = self._get_sizes(

145 + hidden_states.shape[0] // dist_group.world_size,

146 + dist_group,

147 + )

148 hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes)

149 return hidden_states

150

vllm/distributed/device_communicators/base_device_communicator.py

+8 -3

@@ -176,11 +176,16 @@ class DeviceCommunicatorBase:

176 config = get_current_vllm_config_or_none()

177 if config is not None:

178 # initialize the all2all manager for DP or sequence-parallel EP.

179 + parallel_config = config.parallel_config

180 use_ep = (

180 - config.parallel_config.data_parallel_size > 1

181 - or config.parallel_config.use_sequence_parallel_moe

181 + parallel_config.data_parallel_size > 1

182 + or parallel_config.use_sequence_parallel_moe

183 + or (

184 + parallel_config.enable_expert_parallel

185 + and parallel_config.prefill_context_parallel_size > 1

186 + )

187 )

183 - all2all_backend = config.parallel_config.all2all_backend

188 + all2all_backend = parallel_config.all2all_backend

189

190 self.is_ep_communicator = unique_name.split(":")[0] == "ep"

191 self.use_all2all = self.is_ep_communicator and use_ep

vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py

+4 -5

@@ -36,13 +36,12 @@ def build_offloading_config(

36 engine_id = kv_transfer_config.engine_id

38 parallel_config = vllm_config.parallel_config

39 - context_parallel_factor = (

40 - parallel_config.decode_context_parallel_size

41 - * parallel_config.prefill_context_parallel_size

42 - )

39 groups = tuple(

40 OffloadingGroupConfig(

45 - tokens_per_block=(group.kv_cache_spec.block_size * context_parallel_factor),

41 + tokens_per_block=(

42 + group.kv_cache_spec.block_size

43 + * parallel_config.decode_context_parallel_size

44 + ),

45 layer_names=tuple(group.layer_names),

46 )

47 for group in kv_cache_config.kv_cache_groups

vllm/distributed/parallel_state.py

+14 -11

@@ -1784,7 +1784,7 @@ def initialize_model_parallel(

1784 get_world_group().device_group

1785 )

1786

1787 - # the layout order is: ExternalDP x DP x PP x TP

1787 + # the layout order is: ExternalDP x DP x PP x PCP x TP

1788 # ExternalDP is the data parallel group that is not part of the model,

1789 # every dp rank can generate independently (in verl integration).

1790 # DP is the data parallel group that is part of the model,

@@ -1821,17 +1821,13 @@ def initialize_model_parallel(

1821 # Build the DCP model-parallel groups.

1822 global _DCP

1823 assert _DCP is None, "decode context model parallel group is already initialized"

1824 - # Note(hc): In the current implementation of decode context parallel,

1825 - # dcp_size must not exceed tp_size, because the world size does not

1826 - # change by DCP, it simply reuses the GPUs of TP group, and split one

1827 - # TP group into tp_size//dcp_size DCP groups.

1828 - group_ranks = all_ranks.reshape(-1, decode_context_model_parallel_size).unbind(0)

1824 + dcp_size = decode_context_model_parallel_size or 1

1825 + dcp_ranks = local_all_ranks if enable_elastic_ep else all_ranks

1826 + if dcp_size > 1:

1827 + # DCP spans PCP first, then TP for full TP x PCP groups.

1828 + dcp_ranks = dcp_ranks.transpose(-1, -2)

1829 + group_ranks = dcp_ranks.reshape(-1, dcp_size).unbind(0)

1830 group_ranks = [x.tolist() for x in group_ranks]

1830 - if enable_elastic_ep:

1831 - group_ranks = local_all_ranks.reshape(

1832 - -1, decode_context_model_parallel_size

1833 - ).unbind(0)

1834 - group_ranks = [x.tolist() for x in group_ranks]

1831 _DCP = init_model_parallel_group(

1832 group_ranks,

1833 get_world_group().local_rank,

@@ -2005,6 +2001,13 @@ def ensure_model_parallel_initialized(

2001 f"{pcp_world_size=} vs. "

2002 f"{prefill_context_model_parallel_size=}"

2003 )

2004 + dcp_world_size = get_dcp_group().world_size

2005 + dcp_model_parallel_size = decode_context_model_parallel_size or 1

2006 + assert dcp_world_size == dcp_model_parallel_size, (

2007 + "decode context parallel group already initialized, but of unexpected size: "

2008 + f"{dcp_world_size=} vs. "

2009 + f"{dcp_model_parallel_size=}"

2010 + )

2011

2012

2013 def prepare_communication_buffer_for_model(model: torch.nn.Module):

vllm/model_executor/layers/attention/mla_attention.py

+60 -12

@@ -211,6 +211,7 @@ from vllm.config import (

211 from vllm.config.cache import CacheDType

212 from vllm.distributed.parallel_state import (

213 get_dcp_group,

214 + get_tp_group,

215 is_global_first_rank,

216 )

217 from vllm.forward_context import ForwardContext, get_forward_context

@@ -225,6 +226,10 @@ from vllm.model_executor.layers.attention.attention import (

226 from vllm.model_executor.layers.attention.kv_transfer_utils import (

227 maybe_transfer_kv_layer,

228 )

229 +from vllm.model_executor.layers.attention.pcp import (

230 + finalize_mla_pcp_decode,

231 + maybe_gather_mla_latent_cache_inputs,

232 +)

233 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

234 from vllm.model_executor.layers.linear import (

235 ColumnParallelLinear,

@@ -269,7 +274,7 @@ from vllm.v1.attention.backends.utils import (

274 get_dcp_local_seq_lens,

275 split_decodes_and_prefills,

276 )

272 -from vllm.v1.attention.ops.common import cp_lse_ag_out_rs

277 +from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs

278 from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce

279 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states

280 from vllm.v1.attention.selector import get_attn_backend

@@ -499,6 +504,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):

504 self.use_direct_call = not current_platform.opaque_attention_op()

505

506 vllm_config = get_current_vllm_config()

507 + parallel_config = vllm_config.parallel_config

508 + self.use_pcp = parallel_config.prefill_context_parallel_size > 1

509 compilation_config = vllm_config.compilation_config

510 if prefix in compilation_config.static_forward_context:

511 raise ValueError(f"Duplicate layer name: {prefix}")

@@ -611,11 +618,23 @@ class MLAAttention(nn.Module, AttentionLayerBase):

618 assert isinstance(slot_mapping, dict), (

619 f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "

620 )

621 + layer_slot_mapping = slot_mapping.get(self.layer_name)

622 + kv_for_cache, kpe_for_cache, layer_slot_mapping = (

623 + maybe_gather_mla_latent_cache_inputs(

624 + kv_c_normed,

625 + k_pe,

626 + layer_slot_mapping,

627 + attn_metadata.num_decode_tokens

628 + if attn_metadata is not None

629 + else None,

630 + self.use_pcp,

631 + )

632 + )

633 self.impl.do_kv_cache_update( # type: ignore[attr-defined]

615 - kv_c_normed,

616 - k_pe,

634 + kv_for_cache,

635 + kpe_for_cache,

636 self_kv_cache,

618 - slot_mapping.get(self.layer_name),

637 + layer_slot_mapping,

638 self.kv_cache_dtype,

639 self._k_scale,

640 )

@@ -710,6 +729,10 @@ class MLAAttention(nn.Module, AttentionLayerBase):

729 fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype)

730

731 num_actual_toks = attn_metadata.num_actual_tokens

732 + if self.use_pcp and self.impl.dcp_world_size > 1 and quant_key is not None:

733 + raise NotImplementedError(

734 + "MRV2 MLA PCP+DCP does not support fused output quantization yet."

735 + )

736

737 # Inputs and outputs may be padded for CUDA graphs

738 output_padded = output

@@ -833,11 +856,15 @@ class MLAAttention(nn.Module, AttentionLayerBase):

856 else:

857 mqa_q = (mqa_ql_nope, mqa_q_pe)

858 if self.impl.dcp_world_size > 1:

836 - if isinstance(mqa_q, tuple):

837 - # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P)

838 - mqa_q = torch.cat(mqa_q, dim=-1)

839 - # mqa_q do allgather in head dim.

840 - mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)

859 + if self.use_pcp:

860 + if self.impl.dcp_world_size > self.impl.pcp_world_size:

861 + if isinstance(mqa_q, tuple):

862 + mqa_q = torch.cat(mqa_q, dim=-1)

863 + mqa_q = get_tp_group().all_gather(mqa_q, dim=1)

864 + else:

865 + if isinstance(mqa_q, tuple):

866 + mqa_q = torch.cat(mqa_q, dim=-1)

867 + mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)

868

869 # call decode attn

870 if not self.impl.is_sparse:

@@ -846,6 +873,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

873

874 # correct dcp attn_out with lse.

875 if self.impl.dcp_world_size > 1:

876 + assert lse is not None

877 if self.dcp_a2a:

878 attn_out = dcp_a2a_lse_reduce(

879 attn_out,

@@ -853,6 +881,13 @@ class MLAAttention(nn.Module, AttentionLayerBase):

881 get_dcp_group(),

882 is_lse_base_on_e=self.impl.lse_base_on_e,

883 )

884 + elif self.use_pcp:

885 + attn_out = cp_lse_ag_out_ar(

886 + attn_out,

887 + lse,

888 + get_dcp_group(),

889 + is_lse_base_on_e=self.impl.lse_base_on_e,

890 + )

891 else:

892 attn_out = cp_lse_ag_out_rs(

893 attn_out,

@@ -860,6 +895,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):

895 get_dcp_group(),

896 is_lse_base_on_e=self.impl.lse_base_on_e,

897 )

898 + if self.use_pcp:

899 + attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads)

900

901 # v_up projection

902 self._v_up_proj(attn_out, out=mqa_output_slice)

@@ -903,6 +940,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):

940 raise ValueError(f"Unsupported quant_key: {quant_key}")

941 return quant_output

942

943 + if self.use_pcp and output_padded.shape[0] > num_actual_toks:

944 + output_padded[num_actual_toks:].zero_()

945 return output_padded

946

947 def process_weights_after_loading(self, act_dtype: torch.dtype):

@@ -1082,8 +1121,17 @@ def unified_mla_kv_cache_update(

1121 the data dependency between them to ensure torch.compile preserves ordering.

1122 """

1123 layer_name = _resolve_layer_name(layer_name)

1085 - _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)

1124 + attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(

1125 + layer_name

1126 + )

1127 if layer_slot_mapping is not None:

1128 + kv_c_normed, k_pe, layer_slot_mapping = maybe_gather_mla_latent_cache_inputs(

1129 + kv_c_normed,

1130 + k_pe,

1131 + layer_slot_mapping,

1132 + attn_metadata.num_decode_tokens if attn_metadata is not None else None,

1133 + attn_layer.use_pcp,

1134 + )

1135 attn_layer.impl.do_kv_cache_update( # type: ignore[attr-defined]

1136 kv_c_normed,

1137 k_pe,

@@ -1723,6 +1771,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1771 self.compilation_config = vllm_config.compilation_config

1772 self.vllm_config = vllm_config

1773 self.device = device

1774 + self.use_pcp = parallel_config.prefill_context_parallel_size > 1

1775

1776 self.num_heads = self.model_config.get_num_attention_heads(parallel_config)

1777 self.mla_dims = get_mla_dims(self.model_config)

@@ -1735,11 +1784,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1784

1785 try:

1786 self.dcp_world_size = get_dcp_group().world_size

1738 - self.dcp_rank = get_dcp_group().rank_in_group

1787 except AssertionError:

1788 # DCP might not be initialized in testing

1789 self.dcp_world_size = 1

1742 - self.dcp_rank = 0

1790 self.dcp_local_block_size = parallel_config.cp_kv_cache_interleave_size

1791 self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size

1792 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size

@@ -1857,6 +1904,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1904 common_attn_metadata,

1905 decode_threshold=self.reorder_batch_threshold,

1906 require_uniform=(self.query_len_support != QueryLenSupport.VARLEN),

1907 + treat_short_extends_as_decodes=not self.use_pcp,

1908 )

1909 )

1910

vllm/model_executor/layers/attention/pcp.py new

+92

@@ -0,0 +1,92 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import torch

4 +

5 +from vllm.distributed.parallel_state import (

6 + get_pcp_group,

7 + get_tp_group,

8 +)

9 +

10 +

11 +def _gather_prefill_cache_inputs(

12 + tensors: tuple[torch.Tensor, ...],

13 + slot_mapping: torch.Tensor,

14 + num_decode_tokens: int,

15 +) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:

16 + """Keep replicated decode writes local and gather partitioned prefills."""

17 + local_num_tokens = tensors[0].shape[0]

18 + assert all(tensor.shape[0] == local_num_tokens for tensor in tensors)

19 + assert 0 <= num_decode_tokens <= local_num_tokens

20 +

21 + if num_decode_tokens == local_num_tokens:

22 + return tensors, slot_mapping[:num_decode_tokens]

23 +

24 + pcp_group = get_pcp_group()

25 + gathered_prefills = tuple(

26 + pcp_group.all_gather(tensor[num_decode_tokens:].contiguous(), dim=0)

27 + for tensor in tensors

28 + )

29 + pcp_size = pcp_group.world_size

30 + gathered_slot_mapping = slot_mapping[: pcp_size * local_num_tokens]

31 + if num_decode_tokens == 0:

32 + return gathered_prefills, gathered_slot_mapping

33 +

34 + cache_inputs = tuple(

35 + torch.cat((tensor[:num_decode_tokens], gathered_prefill), dim=0)

36 + for tensor, gathered_prefill in zip(tensors, gathered_prefills)

37 + )

38 + rank_slot_mappings = gathered_slot_mapping.view(pcp_size, local_num_tokens)

39 + cache_slot_mapping = torch.cat(

40 + (

41 + rank_slot_mappings[0, :num_decode_tokens],

42 + rank_slot_mappings[:, num_decode_tokens:].flatten(),

43 + )

44 + )

45 + return cache_inputs, cache_slot_mapping

46 +

47 +

48 +def maybe_gather_mla_latent_cache_inputs(

49 + kv_c_normed: torch.Tensor,

50 + k_pe: torch.Tensor,

51 + slot_mapping: torch.Tensor | None,

52 + num_decode_tokens: int | None,

53 + use_pcp: bool,

54 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

55 + if not use_pcp or num_decode_tokens is None:

56 + return kv_c_normed, k_pe, slot_mapping

57 + assert slot_mapping is not None

58 + num_tokens = kv_c_normed.shape[0]

59 + k_pe_flat = k_pe.reshape(num_tokens, -1)

60 + (cache_kv_c, cache_k_pe_flat), cache_slot_mapping = _gather_prefill_cache_inputs(

61 + (kv_c_normed, k_pe_flat),

62 + slot_mapping,

63 + num_decode_tokens,

64 + )

65 + cache_k_pe = cache_k_pe_flat.view(-1, *k_pe.shape[1:])

66 + return cache_kv_c, cache_k_pe, cache_slot_mapping

67 +

68 +

69 +def maybe_gather_indexer_k(

70 + k: torch.Tensor,

71 + slot_mapping: torch.Tensor,

72 + num_decode_tokens: int,

73 + use_pcp: bool,

74 +) -> tuple[torch.Tensor, torch.Tensor]:

75 + if not use_pcp:

76 + return k, slot_mapping

77 + (cache_k,), cache_slot_mapping = _gather_prefill_cache_inputs(

78 + (k,), slot_mapping, num_decode_tokens

79 + )

80 + return cache_k, cache_slot_mapping

81 +

82 +

83 +def finalize_mla_pcp_decode(

84 + output: torch.Tensor,

85 + num_heads: int,

86 +) -> torch.Tensor:

87 + if output.shape[1] < num_heads:

88 + output = get_pcp_group().all_gather(output, dim=1)

89 + elif output.shape[1] > num_heads:

90 + head_start = get_tp_group().rank_in_group * num_heads

91 + output = output[:, head_start : head_start + num_heads]

92 + return output

vllm/model_executor/layers/attention/sparse_mla_attention.py

+2 -2

@@ -59,13 +59,12 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):

59 device=device,

60 )

61 parallel_config = vllm_config.parallel_config

62 + self.use_pcp = parallel_config.prefill_context_parallel_size > 1

63 try:

64 self.dcp_world_size = get_dcp_group().world_size

64 - self.dcp_rank = get_dcp_group().rank_in_group

65 except AssertionError:

66 # DCP might not be initialized in testing

67 self.dcp_world_size = 1

68 - self.dcp_rank = 0

68 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size

69 self.dcp_local_block_size = self.cp_kv_cache_interleave_size

70 self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size

@@ -174,6 +173,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):

173 num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills(

174 common_attn_metadata,

175 decode_threshold=self.reorder_batch_threshold or 1,

176 + treat_short_extends_as_decodes=not self.use_pcp,

177 require_uniform=self.require_uniform_decodes,

178 )

179 (

vllm/model_executor/layers/fused_moe/config.py

+3 -1

@@ -1050,7 +1050,9 @@ class FusedMoEParallelConfig:

1050

1051 @property

1052 def use_all2all_kernels(self):

1053 - return self.use_ep and (self.dp_size > 1 or self.is_sequence_parallel)

1053 + return self.use_ep and (

1054 + self.dp_size > 1 or self.pcp_size > 1 or self.is_sequence_parallel

1055 + )

1056

1057 @property

1058 def use_deepep_ht_kernels(self):

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

+11 -17

@@ -752,18 +752,12 @@ class MoERunner(MoERunnerInterface):

752 assert len(result) == 2

753 hidden_states, router_logits = result

754

755 - # NOTE: Similar with DP, PCP also needs dispatch and combine. For

756 - # simplicity, AgRsAll2All was added separately for PCP here. Maybe

757 - # we should modify All2AllManager abstraction to better support PCP.

758 - if self.moe_config.pcp_size > 1:

759 - hidden_states = get_pcp_group().all_gather(

760 - hidden_states,

761 - dim=0,

762 - )

763 - router_logits = get_pcp_group().all_gather(

764 - router_logits,

765 - dim=0,

766 - )

755 + if (

756 + self.moe_config.pcp_size > 1

757 + and not self.moe_config.moe_parallel_config.use_all2all_kernels

758 + ):

759 + hidden_states = get_pcp_group().all_gather(hidden_states, dim=0)

760 + router_logits = get_pcp_group().all_gather(router_logits, dim=0)

761

762 return hidden_states, router_logits

763

@@ -777,11 +771,11 @@ class MoERunner(MoERunnerInterface):

771 hidden_states, self.moe_config.is_sequence_parallel

772 )

773

780 - if self.moe_config.pcp_size > 1:

781 - hidden_states = get_pcp_group().reduce_scatter(

782 - hidden_states,

783 - dim=0,

784 - )

774 + if (

775 + self.moe_config.pcp_size > 1

776 + and not self.moe_config.moe_parallel_config.use_all2all_kernels

777 + ):

778 + hidden_states = get_pcp_group().reduce_scatter(hidden_states, dim=0)

779

780 if self.shared_experts is not None:

781 assert shared_output is not None

vllm/model_executor/layers/sparse_attn_indexer.py

+26 -3

@@ -9,10 +9,11 @@ from vllm import _custom_ops as ops

9 from vllm._aiter_ops import rocm_aiter_ops

10 from vllm.compilation.breakable_cudagraph import eager_break_during_capture

11 from vllm.config import get_current_vllm_config

12 -from vllm.distributed import get_dcp_group

12 +from vllm.distributed import get_dcp_group, get_pcp_group

13 from vllm.forward_context import get_forward_context

14 from vllm.logger import init_logger

15 from vllm.model_executor.custom_op import CustomOp

16 +from vllm.model_executor.layers.attention.pcp import maybe_gather_indexer_k

17 from vllm.model_executor.layers.quantization.utils.quant_utils import (

18 get_fp8_min_max,

19 )

@@ -308,6 +309,7 @@ def sparse_attn_indexer(

309 total_seq_lens: int,

310 topk_indices_buffer: torch.Tensor,

311 skip_k_cache_insert: bool,

312 + use_pcp: bool,

313 use_fp4_cache: bool = False,

314 dcp_rank: int = 0,

315 dcp_world_size: int = 1,

@@ -354,6 +356,7 @@ def sparse_attn_indexer(

356 total_seq_lens,

357 topk_indices_buffer,

358 skip_k_cache_insert,

359 + use_pcp,

360 use_fp4_cache,

361 )

362 attn_metadata_narrowed = attn_metadata[k_cache_prefix]

@@ -373,18 +376,28 @@ def sparse_attn_indexer(

376 # During speculative decoding, k may be padded to the CUDA graph batch

377 # size while slot_mapping only covers actual tokens. Truncate k to avoid

378 # out-of-bounds reads in the kernel.

379 + # Keep PCP padding so every rank contributes the same all-gather shape.

380 num_tokens = slot_mapping.shape[0]

381 + if use_pcp:

382 + num_tokens //= get_pcp_group().world_size

383 if k is not None:

384 k = k[:num_tokens]

385

386 if not skip_k_cache_insert:

387 + assert k is not None

388 + k, slot_mapping_for_cache = maybe_gather_indexer_k(

389 + k,

390 + slot_mapping,

391 + num_decode_tokens,

392 + use_pcp,

393 + )

394 # scale_fmt can be None, but the function expects str

395 assert scale_fmt is not None

396 assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet"

397 ops.indexer_k_quant_and_cache(

398 k,

399 kv_cache,

387 - slot_mapping,

400 + slot_mapping_for_cache,

401 quant_block_size,

402 scale_fmt,

403 )

@@ -498,7 +511,14 @@ def sparse_attn_indexer(

511 assert decode_metadata is not None

512 kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)

513 decode_lens = decode_metadata.decode_lens

501 - if decode_metadata.requires_padding:

514 + if num_decode_tokens == 0:

515 + padded_q_quant_decode_tokens = q_quant[:1].reshape(1, 1, *q_quant.shape[1:])

516 + padded_q_scale = (

517 + q_scale[:1].reshape(1, 1, *q_scale.shape[1:])

518 + if q_scale is not None

519 + else None

520 + )

521 + elif decode_metadata.requires_padding:

522 # pad in edge case where we have short chunked prefill length <

523 # decode_threshold since we unstrictly split

524 # prefill and decode by decode_threshold

@@ -663,6 +683,7 @@ def sparse_attn_indexer_fake(

683 total_seq_lens: int,

684 topk_indices_buffer: torch.Tensor | None,

685 skip_k_cache_insert: bool,

686 + use_pcp: bool,

687 use_fp4_cache: bool = False,

688 dcp_rank: int = 0,

689 dcp_world_size: int = 1,

@@ -725,6 +746,7 @@ class SparseAttnIndexer(CustomOp):

746 self.dcp_world_size = parallel_config.decode_context_parallel_size

747 self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0

748 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size

749 + self.use_pcp = parallel_config.prefill_context_parallel_size > 1

750 if current_platform.is_cuda() and not has_deep_gemm():

751 raise RuntimeError(

752 "Sparse Attention Indexer CUDA op requires DeepGEMM support in "

@@ -777,6 +799,7 @@ class SparseAttnIndexer(CustomOp):

799 self.max_total_seq_len,

800 self.topk_indices_buffer,

801 self.skip_k_cache_insert,

802 + self.use_pcp,

803 self.use_fp4_cache,

804 self.dcp_rank,

805 self.dcp_world_size,

vllm/v1/attention/backend.py

+14 -2

@@ -281,6 +281,13 @@ class AttentionBackend(ABC):

281 def supports_kv_connector(cls) -> bool:

282 return True

283

284 + @classmethod

285 + def supports_pcp(cls) -> bool:

286 + try:

287 + return cls.get_impl_cls().supports_pcp

288 + except NotImplementedError:

289 + return False

290 +

291 @classmethod

292 def supports_attn_type(cls, attn_type: str) -> bool:

293 """Check if backend supports a given attention type.

@@ -327,6 +334,7 @@ class AttentionBackend(ABC):

334 use_non_causal: bool = False,

335 use_batch_invariant: bool = False,

336 use_kv_connector: bool = False,

337 + use_pcp: bool = False,

338 ) -> list[str]:

339 invalid_reasons = []

340 if not cls.supports_head_size(head_size):

@@ -367,6 +375,8 @@ class AttentionBackend(ABC):

375 invalid_reasons.append("batch invariance not supported")

376 if use_kv_connector and not cls.supports_kv_connector():

377 invalid_reasons.append("KV connector not supported")

378 + if use_pcp and not cls.supports_pcp():

379 + invalid_reasons.append("PCP not supported")

380 combination_reason = cls.supports_combination(

381 head_size,

382 dtype,

@@ -858,8 +868,8 @@ class AttentionImplBase(ABC, Generic[T]):

868 except AssertionError:

869 self.pcp_world_size = 1

870 self.pcp_rank = 0

861 - self.total_cp_world_size = self.pcp_world_size * self.dcp_world_size

862 - self.total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank

871 + self.total_cp_world_size = self.dcp_world_size

872 + self.total_cp_rank = self.dcp_rank

873

874 self.need_to_return_lse_for_decode = (

875 self.dcp_world_size > 1 and self.can_return_lse_for_decode

@@ -987,6 +997,8 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):

997 class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):

998 """MLA attention implementation with forward_mqa and forward_mha methods."""

999

1000 + supports_pcp: bool = True

1001 +

1002 @abstractmethod

1003 def __init__(

1004 self,

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

+19 -4

@@ -6,7 +6,7 @@ import torch

7 import vllm.envs as envs

8 from vllm.config import VllmConfig

9 -from vllm.distributed import get_dcp_group

9 +from vllm.distributed import get_dcp_group, get_pcp_group

10 from vllm.logger import init_logger

11 from vllm.platforms import current_platform

12 from vllm.triton_utils import tl, triton

@@ -29,7 +29,7 @@ from vllm.v1.attention.backends.utils import (

29 split_decodes_and_prefills,

30 )

31 from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec

32 -from vllm.v1.worker.cp_utils import get_total_cp_world_size

32 +from vllm.v1.worker.cp_utils import get_kv_cache_shard_count

34 logger = init_logger(__name__)

@@ -109,7 +109,7 @@ def split_indexer_prefill_chunks(

109 end += 1

110

111 req_slice = slice(start + request_offset, end + request_offset)

112 - max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else chunk_m

112 + max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m)

113 for q_off in range(0, chunk_m, max_q):

114 sub_m = min(max_q, chunk_m - q_off)

115 chunks.append((req_slice, slice(q_off, q_off + sub_m)))

@@ -118,6 +118,10 @@ def split_indexer_prefill_chunks(

118

119

120 class DeepseekV32IndexerBackend(AttentionBackend):

121 + @classmethod

122 + def supports_pcp(cls) -> bool:

123 + return True

124 +

125 @staticmethod

126 def get_name() -> str:

127 return "DEEPSEEK_V32_INDEXER"

@@ -256,6 +260,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

260 parallel_config = self.vllm_config.parallel_config

261 self.dcp_world_size = parallel_config.decode_context_parallel_size

262 self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0

263 + self.pcp_world_size = parallel_config.prefill_context_parallel_size

264 + self.use_pcp = self.pcp_world_size > 1

265 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size

266 # The DCP sparse-indexer code is parameterized by interleave size, but

267 # interleave > 1 is not yet validated end-to-end (gsm8k parity fails),

@@ -339,7 +345,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

345 )

346 max_num_blocks_per_req = cdiv(

347 self.vllm_config.model_config.max_model_len,

342 - self.kv_cache_spec.block_size * get_total_cp_world_size(),

348 + self.kv_cache_spec.block_size * get_kv_cache_shard_count(),

349 )

350 self.expanded_block_table_buffer = torch.zeros(

351 (

@@ -570,6 +576,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

576 common_attn_metadata,

577 decode_threshold=self.decode_threshold,

578 require_uniform=not self.use_flattening,

579 + treat_short_extends_as_decodes=not self.use_pcp,

580 )

581 )

582

@@ -579,6 +586,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

586 compressed_slot_mapping = slot_mapping

587 compressed_seq_lens = seq_lens

588 if self.compress_ratio > 1:

589 + padded_num_tokens = num_tokens

590 + if self.pcp_world_size > 1:

591 + padded_num_tokens = slot_mapping.shape[0] // self.pcp_world_size

592 compressed_slot_mapping = get_compressed_slot_mapping(

593 num_tokens,

594 query_start_loc,

@@ -588,6 +598,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

598 self.compress_ratio,

599 out=self.compressed_slot_mapping_buffer,

600 )

601 + if self.pcp_world_size > 1:

602 + compressed_slot_mapping = get_pcp_group().all_gather(

603 + self.compressed_slot_mapping_buffer[:padded_num_tokens],

604 + dim=0,

605 + )

606 compressed_seq_lens = seq_lens // self.compress_ratio

607

608 prefill_metadata = None

vllm/v1/attention/backends/utils.py

+3 -1

@@ -612,7 +612,9 @@ def split_decodes_and_prefills(

612 # check if we are in a padded uniform batch; this is used for full-CGs, some

613 # requests may have a query length of 0 but since they are padding its fine

614 # to treat them as decodes (ensures num_decodes matches the captured size)

615 - if torch.all((query_lens == query_lens[0]) | (query_lens == 0)):

615 + if treat_short_extends_as_decodes and torch.all(

616 + (query_lens == query_lens[0]) | (query_lens == 0)

617 + ):

618 return num_reqs, 0, num_tokens, 0 # all decodes

619 is_prefill = query_lens != query_lens[0]

620 else:

vllm/v1/attention/selector.py

+4 -1

@@ -36,6 +36,7 @@ class AttentionSelectorConfig(NamedTuple):

36 use_non_causal: bool = False

37 use_batch_invariant: bool = False

38 use_kv_connector: bool = False

39 + use_pcp: bool = False

41 def __repr__(self):

42 return (

@@ -52,7 +53,8 @@ class AttentionSelectorConfig(NamedTuple):

53 f"has_sliding_window={self.has_sliding_window}, "

54 f"use_non_causal={self.use_non_causal}, "

55 f"use_batch_invariant={self.use_batch_invariant}, "

55 - f"use_kv_connector={self.use_kv_connector})"

56 + f"use_kv_connector={self.use_kv_connector}, "

57 + f"use_pcp={self.use_pcp})"

58 )

@@ -150,6 +152,7 @@ def get_attn_backend(

152 use_non_causal=vllm_config.attention_config.use_non_causal,

153 use_batch_invariant=envs.VLLM_BATCH_INVARIANT,

154 use_kv_connector=use_kv_connector,

155 + use_pcp=vllm_config.parallel_config.prefill_context_parallel_size > 1,

156 )

157

158 # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over

vllm/v1/core/kv_cache_coordinator.py

-2

@@ -472,8 +472,6 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

472 self.pcp_world_size = pcp_world_size

473 if dcp_world_size > 1:

474 self.block_size *= dcp_world_size

475 - if pcp_world_size > 1:

476 - self.block_size *= pcp_world_size

475 # For models using only Mamba, block_size is set to max_model_len when

476 # prefix caching is disabled, and hash_block_size validation is skipped.

477 assert not enable_caching or (hash_block_size == self.block_size), (

vllm/v1/core/kv_cache_utils.py

+5 -6

@@ -631,8 +631,8 @@ def resolve_kv_cache_block_sizes(

631

632 - ``scheduler_block_size`` is the token-alignment invariant used by the

633 scheduler (e.g. for ``num_computed_tokens`` rounding). Single group:

634 - ``cache_config.block_size * dcp * pcp``. Multiple groups: LCM of every

635 - group's effective block size. Attention groups are scaled by DCP/PCP;

634 + ``cache_config.block_size * dcp``. Multiple groups: LCM of every

635 + group's effective block size. Attention groups are scaled by DCP;

636 Mamba groups keep their full per-rank state and are not scaled.

637 - ``hash_block_size`` is the granularity at which ``Request.block_hashes``

638 is computed. Single group: equals scheduler block size. Multiple groups:

@@ -644,15 +644,14 @@ def resolve_kv_cache_block_sizes(

644 """

645 cache_config = vllm_config.cache_config

646 dcp = vllm_config.parallel_config.decode_context_parallel_size

647 - pcp = vllm_config.parallel_config.prefill_context_parallel_size

647 groups = kv_cache_config.kv_cache_groups

648

650 - if len(groups) <= 1: # Single group: block_size * dcp * pcp

651 - bs = cache_config.block_size * dcp * pcp

649 + if len(groups) <= 1:

650 + bs = cache_config.block_size * dcp

651 return bs, bs

652

653 group_block_sizes = [

655 - g.kv_cache_spec.block_size * dcp * pcp

654 + g.kv_cache_spec.block_size * dcp

655 if isinstance(g.kv_cache_spec, AttentionSpec)

656 else g.kv_cache_spec.block_size

657 for g in groups

vllm/v1/core/sched/scheduler.py

+1 -1

@@ -271,7 +271,7 @@ class Scheduler(SchedulerInterface):

271 log_stats=self.log_stats,

272 enable_kv_cache_events=self.enable_kv_cache_events,

273 dcp_world_size=self.dcp_world_size,

274 - pcp_world_size=self.pcp_world_size,

274 + pcp_world_size=1,

275 scheduler_block_size=self.block_size,

276 hash_block_size=hash_block_size,

277 metrics_collector=self.kv_metrics_collector,

vllm/v1/core/single_type_kv_cache_manager.py

+6 -6

@@ -75,8 +75,8 @@ class SingleTypeKVCacheManager(ABC):

75 self.block_size = kv_cache_spec.block_size

76 self.dcp_world_size = dcp_world_size

77 self.pcp_world_size = pcp_world_size

78 - if dcp_world_size * pcp_world_size > 1:

79 - self.block_size *= dcp_world_size * pcp_world_size

78 + if dcp_world_size > 1:

79 + self.block_size *= dcp_world_size

80 self.kv_cache_spec = kv_cache_spec

81 self.block_pool = block_pool

82 self.enable_caching = enable_caching

@@ -674,10 +674,10 @@ class FullAttentionManager(SingleTypeKVCacheManager):

674 "and chunked local attention groups"

675 )

676 block_size = kv_cache_spec.block_size

677 - if dcp_world_size * pcp_world_size > 1:

678 - # DCP/PCP shard each block's KV across ranks; hashes must be

679 - # viewed at the sharded (scaled) block size.

680 - block_size *= dcp_world_size * pcp_world_size

677 + if dcp_world_size > 1:

678 + # DCP shards each block's KV across ranks; hashes must be viewed at

679 + # the sharded block size.

680 + block_size *= dcp_world_size

681 block_hashes = resolve_block_hashes(

682 block_hashes,

683 block_pool.hash_block_size,

vllm/v1/kv_cache_interface.py

+4 -12

@@ -218,14 +218,9 @@ class AttentionSpec(KVCacheSpec):

218 )

219

220 def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:

221 - # Attention KV is token-interleaved across DCP/PCP ranks, so each rank

222 - # only stores max_len // (dcp * pcp) tokens per request.

221 parallel_config = vllm_config.parallel_config

224 - total_cp_size = (

225 - parallel_config.decode_context_parallel_size

226 - * parallel_config.prefill_context_parallel_size

227 - )

228 - return cdiv(max_len, self.block_size * total_cp_size)

222 + kv_shard_count = parallel_config.decode_context_parallel_size

223 + return cdiv(max_len, self.block_size * kv_shard_count)

224

225

226 @dataclass(frozen=True, kw_only=True)

@@ -263,11 +258,8 @@ class FullAttentionSpec(AttentionSpec):

258 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

259 max_model_len = vllm_config.model_config.max_model_len

260 dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size

266 - pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size

267 - # Note(hc): each dcp rank only need save

268 - # (max_model_len//dcp_world_size) tokens locally.

269 - if dcp_world_size * pcp_world_size > 1:

270 - max_model_len = cdiv(max_model_len, dcp_world_size * pcp_world_size)

261 + if dcp_world_size > 1:

262 + max_model_len = cdiv(max_model_len, dcp_world_size)

263 return cdiv(max_model_len, self.block_size) * self.page_size_bytes

264

265 @classmethod

vllm/v1/simple_kv_offload/manager.py

+3 -4

@@ -83,8 +83,7 @@ class SimpleCPUOffloadScheduler:

83 and vllm_config.kv_events_config.enable_kv_cache_events

84 )

85 dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size

86 - pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size

87 - self.cp_world_size = dcp_world_size * pcp_world_size

86 + self.cp_world_size = dcp_world_size

87 self.block_size = scheduler_block_size

88 self.hash_block_size = hash_block_size

89 assert self.block_size % self.hash_block_size == 0

@@ -127,7 +126,7 @@ class SimpleCPUOffloadScheduler:

126 enable_caching=True,

127 enable_kv_cache_events=self.enable_kv_cache_events,

128 dcp_world_size=dcp_world_size,

130 - pcp_world_size=pcp_world_size,

129 + pcp_world_size=1,

130 scheduler_block_size=self.block_size,

131 hash_block_size=self.hash_block_size,

132 )

@@ -354,7 +353,7 @@ class SimpleCPUOffloadScheduler:

353 )

354 assert num_external_tokens % g_block_size == 0, (

355 f"num_external_tokens={num_external_tokens} not aligned to "

357 - f"group {g} effective block_size={g_block_size}"

356 + f"group {g} block_size={g_block_size}"

357 )

358 n_take_g = num_external_tokens // g_block_size

359 cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g])

vllm/v1/worker/block_table.py

+2 -4

@@ -163,8 +163,6 @@ class BlockTable:

163 return

164 assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT

165

166 - total_cp_world_size = self.pcp_world_size * self.dcp_world_size

167 - total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank

166 _compute_slot_mapping_kernel[(num_reqs + 1,)](

167 num_tokens,

168 self.max_num_batched_tokens,

@@ -176,8 +174,8 @@ class BlockTable:

174 self.slot_mapping.gpu,

175 KV_CACHE_BLOCK_SIZE=self.kv_cache_block_size,

176 BLOCKS_PER_KV_BLOCK=self.blocks_per_kv_block,

179 - TOTAL_CP_WORLD_SIZE=total_cp_world_size,

180 - TOTAL_CP_RANK=total_cp_rank,

177 + TOTAL_CP_WORLD_SIZE=self.dcp_world_size,

178 + TOTAL_CP_RANK=self.dcp_rank,

179 CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size,

180 PAD_ID=PAD_SLOT_ID,

181 BLOCK_SIZE=1024,

vllm/v1/worker/cp_utils.py

+10 -15

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, cast

6 import torch

8 from vllm.config import VllmConfig, get_layers_from_vllm_config

9 -from vllm.distributed import get_dcp_group, get_pcp_group

9 +from vllm.distributed import get_dcp_group

10 from vllm.logger import init_logger

11 from vllm.v1.attention.backend import CommonAttentionMetadata

12 from vllm.v1.attention.backends.utils import split_decodes_prefills_and_extends

@@ -27,6 +27,13 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:

27 layer_type = cast(type[Any], AttentionLayerBase)

28 layers = get_layers_from_vllm_config(vllm_config, layer_type)

29 for layer in layers.values():

30 + get_attn_backend = getattr(layer, "get_attn_backend", None)

31 + if pcp_size > 1 and get_attn_backend is not None:

32 + backend = get_attn_backend()

33 + assert backend.supports_pcp(), (

34 + "PCP requires attention backend support, "

35 + f"but {backend.get_name()} does not support PCP."

36 + )

37 layer_impl = getattr(layer, "impl", None)

38 if layer_impl is None:

39 continue

@@ -44,26 +51,14 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:

51 "--attention-backend or disable DCP."

52 )

47 - if pcp_size > 1:

48 - assert layer_impl.supports_pcp, (

49 - "PCP requires attention impls' support, "

50 - f"but the impl {layer_impl.__class__.__name__} "

51 - "does not support PCP."

52 - )

53 -

55 -def get_total_cp_world_size():

56 - try:

57 - pcp_world_size = get_pcp_group().world_size

58 - except AssertionError:

59 - # PCP might not be initialized in testing

60 - pcp_world_size = 1

55 +def get_kv_cache_shard_count() -> int:

56 try:

57 dcp_world_size = get_dcp_group().world_size

58 except AssertionError:

59 # DCP might not be initialized in testing

60 dcp_world_size = 1

66 - return dcp_world_size * pcp_world_size

61 + return dcp_world_size

64 def get_dcp_dummy_context_len(

vllm/v1/worker/gpu/attn_utils.py

@@ -578,6 +578,7 @@ def build_attn_metadata(

578 seq_lens_cpu_upper_bound: torch.Tensor | None = None,

579 dcp_local_seq_lens: torch.Tensor | None = None,

580 positions: torch.Tensor | None = None,

581 + is_prefilling: torch.Tensor | None = None,

582 mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None,

583 model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,

584 for_cudagraph_capture: bool = False,

@@ -605,6 +606,11 @@ def build_attn_metadata(

606 if model_specific_attn_metadata is not None

607 else {}

608 )

609 + # Model-specific metadata (e.g. Mamba hybrid) may supply its own

610 + # padding-aware is_prefilling, which takes precedence over the default.

611 + group_is_prefilling = common_attn_metadata_extra_kwargs.pop(

612 + "is_prefilling", is_prefilling

613 + )

614 common_attn_metadata = CommonAttentionMetadata(

615 query_start_loc=query_start_loc_gpu,

616 query_start_loc_cpu=query_start_loc_cpu,

@@ -619,6 +625,7 @@ def build_attn_metadata(

625 causal=group_causal,

626 dcp_local_seq_lens=dcp_local_seq_lens,

627 positions=positions,

628 + is_prefilling=group_is_prefilling,

629 mm_req_doc_ranges=mm_req_doc_ranges,

630 rswa_prefix_lens=rswa_prefix_lens,

631 **common_attn_metadata_extra_kwargs,

vllm/v1/worker/gpu/block_table.py

+16 -6

@@ -135,20 +135,28 @@ class BlockTables:

135 self,

136 idx_mapping: torch.Tensor,

137 num_reqs_padded: int,

138 + out: tuple[torch.Tensor, ...] | None = None,

139 + out_ptrs: torch.Tensor | None = None,

140 ) -> tuple[torch.Tensor, ...]:

141 + if out is None:

142 + out = tuple(self.input_block_tables)

143 + out_ptrs = self.input_block_table_ptrs

144 + else:

145 + assert out_ptrs is not None

146 + assert len(out) == self.num_kv_cache_groups

147 num_reqs = idx_mapping.shape[0]

148 # Launch kernel with num_reqs_padded to fuse zeroing of padded rows.

149 _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs_padded)](

150 idx_mapping,

151 self.block_table_ptrs,

144 - self.input_block_table_ptrs,

152 + out_ptrs,

153 self.block_table_strides,

154 self.num_blocks.gpu,

155 self.num_blocks.gpu.stride(0),

156 num_reqs,

157 BLOCK_SIZE=1024, # type: ignore

158 )

151 - return tuple(bt[:num_reqs_padded] for bt in self.input_block_tables)

159 + return tuple(bt[:num_reqs_padded] for bt in out)

160

161 def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:

162 # NOTE(woosuk): The output may be used for CUDA graph capture.

@@ -163,26 +171,28 @@ class BlockTables:

171 query_start_loc: torch.Tensor,

172 positions: torch.Tensor,

173 num_tokens_padded: int,

174 + out: torch.Tensor | None = None,

175 ) -> torch.Tensor:

176 num_reqs = idx_mapping.shape[0]

177 num_groups = self.num_kv_cache_groups

178 + slot_mappings = self.slot_mappings if out is None else out

179 _compute_slot_mappings_kernel[(num_groups, num_reqs + 1)](

170 - self.max_num_batched_tokens,

180 + slot_mappings.shape[1],

181 idx_mapping,

182 query_start_loc,

183 positions,

184 self.block_table_ptrs,

185 self.block_table_strides,

186 self.block_sizes_tensor,

177 - self.slot_mappings,

178 - self.slot_mappings.stride(0),

187 + slot_mappings,

188 + slot_mappings.stride(0),

189 self.cp_rank,

190 CP_SIZE=self.cp_size,

191 CP_INTERLEAVE=self.cp_interleave,

192 PAD_ID=PAD_SLOT_ID,

193 TRITON_BLOCK_SIZE=1024, # type: ignore

194 )

185 - return self.slot_mappings[:, :num_tokens_padded]

195 + return slot_mappings[:, :num_tokens_padded]

196

197 def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:

198 # Fill the entire slot_mappings tensor, not just the first `num_tokens` entries.

vllm/v1/worker/gpu/model_runner.py

+21 -3

@@ -52,6 +52,7 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput

52 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec

53 from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput

54 from vllm.v1.worker.cp_utils import check_attention_cp_compatibility

55 +from vllm.v1.worker.gpu import pcp_manager as pcp

56 from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput

57 from vllm.v1.worker.gpu.attn_utils import (

58 build_slot_mappings_by_layer,

@@ -205,6 +206,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

206 # Draft tokens propagation - for spec-dec + struct outputs.

207 self.draft_tokens_handler = DraftTokensHandler(self.device)

208

209 + self.pcp_manager: pcp.PCPManager | None = None

210 +

211 # Pooling models.

212 self.is_pooling_model = self.model_config.runner_type == "pooling"

213 self.pooling_runner: PoolingRunner | None = None

@@ -456,6 +459,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):

459 cp_rank=self.dcp_rank,

460 cp_interleave=self.cp_interleave,

461 )

462 + self.pcp_manager = pcp.maybe_build_pcp_manager(

463 + self.vllm_config,

464 + self.device,

465 + self.supports_mm_inputs,

466 + self.req_states,

467 + self.block_tables,

468 + )

469 initialize_mamba_ssu_backend(

470 self.vllm_config.mamba_config, self.kv_cache_config

471 )

@@ -1008,7 +1018,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1018 # prompt_lens is only used in R-SWA case.

1019 prompt_lens = self.req_states.prompt_len.gpu[idx_mapping]

1020

1011 - return InputBatch(

1021 + input_batch = InputBatch(

1022 req_ids=req_ids,

1023 num_reqs=num_reqs,

1024 num_reqs_after_padding=num_reqs_padded,

@@ -1040,10 +1050,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1050 has_structured_output_reqs=scheduler_output.has_structured_output_requests,

1051 prompt_lens=prompt_lens,

1052 )

1053 + return pcp.maybe_partition_pcp_batch(self.pcp_manager, input_batch)

1054

1055 def prepare_attn(

1056 self, input_batch: InputBatch

1057 ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:

1058 + if self.pcp_manager is not None:

1059 + return self.pcp_manager.prepare_attn(input_batch)

1060 +

1061 # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].

1062 block_tables = self.block_tables.gather_block_tables(

1063 input_batch.idx_mapping,

@@ -1063,8 +1077,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1077 self, input_batch: InputBatch

1078 ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:

1079 block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)

1066 - slot_mappings = self.block_tables.get_dummy_slot_mappings(

1067 - input_batch.num_tokens

1080 + slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings(

1081 + self.pcp_manager, self.block_tables, input_batch.num_tokens

1082 )

1083 return block_tables, slot_mappings

1084

@@ -1412,6 +1426,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1426 return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)

1427

1428 # Last rank: sample tokens

1429 + hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling(

1430 + self.pcp_manager, hidden_states, input_batch

1431 + )

1432 +

1433 sampler_output, num_sampled, num_rejected = self.sample(

1434 hidden_states, input_batch, grammar_output

1435 )

vllm/v1/worker/gpu/model_states/default.py

+6 -2

@@ -151,7 +151,10 @@ class DefaultModelState(ModelState):

151 # For piecewise cudagraphs and eager, use unpadded sizes.

152 num_reqs = input_batch.num_reqs

153 num_tokens = input_batch.num_tokens

154 - query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)

154 + query_start_loc_cpu = torch.from_numpy(

155 + input_batch.query_start_loc_np[: num_reqs + 1]

156 + )

157 + query_start_loc_gpu = input_batch.query_start_loc[: num_reqs + 1]

158 max_query_len = input_batch.num_scheduled_tokens.max().item()

159 seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound

160 if for_capture:

@@ -174,7 +177,7 @@ class DefaultModelState(ModelState):

177 attn_groups=attn_groups,

178 num_reqs=num_reqs,

179 num_tokens=num_tokens,

177 - query_start_loc_gpu=input_batch.query_start_loc,

180 + query_start_loc_gpu=query_start_loc_gpu,

181 query_start_loc_cpu=query_start_loc_cpu,

182 max_query_len=max_query_len,

183 seq_lens=input_batch.seq_lens,

@@ -185,6 +188,7 @@ class DefaultModelState(ModelState):

188 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

189 dcp_local_seq_lens=input_batch.dcp_local_seq_lens,

190 positions=input_batch.positions,

191 + is_prefilling=torch.from_numpy(input_batch.is_prefilling_np),

192 mm_req_doc_ranges=req_doc_ranges,

193 for_cudagraph_capture=for_capture,

194 rswa_prefix_lens=input_batch.prompt_lens,

vllm/v1/worker/gpu/pcp_manager.py new

+680

@@ -0,0 +1,680 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from dataclasses import dataclass, replace

4 +

5 +import numpy as np

6 +import torch

7 +

8 +from vllm.config import CUDAGraphMode, VllmConfig

9 +from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group

10 +from vllm.logger import init_logger

11 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID

12 +from vllm.v1.worker.gpu.block_table import BlockTables

13 +from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu

14 +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens

15 +from vllm.v1.worker.gpu.input_batch import (

16 + InputBatch,

17 + InputBuffers,

18 + combine_sampled_and_draft_tokens,

19 + prepare_pos_seq_lens,

20 +)

21 +from vllm.v1.worker.gpu.states import RequestState

22 +

23 +logger = init_logger(__name__)

24 +

25 +

26 +@dataclass(frozen=True)

27 +class RankSegment:

28 + global_batch_req_idx: int

29 + global_batch_slice: slice

30 + rank_local_batch_slice: slice

31 +

32 + @property

33 + def num_tokens(self) -> int:

34 + return self.global_batch_slice.stop - self.global_batch_slice.start

35 +

36 +

37 +class PCPManager:

38 + """MRV2 PC batch manager.

39 +

40 + The model runner keeps the global scheduled batch. This manager rewrites only

41 + the per-step InputBatch into rank-local DualChunkSwap rows and keeps the

42 + global-batch view private to restore to the global batch shape before

43 + sampling/postprocess.

44 + """

45 +

46 + def __init__(

47 + self,

48 + pcp_world_size: int,

49 + pcp_rank: int,

50 + device: torch.device,

51 + req_states: RequestState | None = None,

52 + max_num_reqs: int | None = None,

53 + max_num_tokens: int | None = None,

54 + block_tables: BlockTables | None = None,

55 + dcp_world_size: int = 1,

56 + dcp_rank: int = 0,

57 + cp_interleave: int = 1,

58 + ) -> None:

59 + self.pcp_world_size = pcp_world_size

60 + self.pcp_rank = pcp_rank

61 + self.device = device

62 + self.dcp_world_size = dcp_world_size

63 + self.dcp_rank = dcp_rank

64 + self.cp_interleave = cp_interleave

65 +

66 + self._global_batch: InputBatch | None = None

67 + self._req_states = req_states

68 + self._block_tables = block_tables

69 + self._hidden_restore_idx: torch.Tensor | None = None

70 + self._padded_gather_idx: torch.Tensor | None = None

71 + self._gathered_kv_write_mask: torch.Tensor | None = None

72 + self._pad_slot_id = torch.tensor(PAD_SLOT_ID, dtype=torch.int64, device=device)

73 +

74 + max_num_local_reqs = 2 * max_num_reqs if max_num_reqs is not None else None

75 + self._input_buffers = (

76 + InputBuffers(max_num_local_reqs, max_num_tokens, device)

77 + if max_num_local_reqs is not None and max_num_tokens is not None

78 + else None

79 + )

80 + self._local_req_idx = (

81 + torch.arange(max_num_local_reqs, dtype=torch.int32, device=device)

82 + if max_num_local_reqs is not None

83 + else None

84 + )

85 + self._local_block_tables: tuple[torch.Tensor, ...] | None

86 + self._local_block_table_ptrs: torch.Tensor | None

87 + if block_tables is not None and max_num_local_reqs is not None:

88 + self._local_block_tables = tuple(

89 + table.new_zeros((max_num_local_reqs, table.shape[1]))

90 + for table in block_tables.input_block_tables

91 + )

92 + self._local_block_table_ptrs = torch.tensor(

93 + [table.data_ptr() for table in self._local_block_tables],

94 + dtype=torch.uint64,

95 + device=device,

96 + )

97 + else:

98 + self._local_block_tables = None

99 + self._local_block_table_ptrs = None

100 + num_kv_cache_groups = (

101 + block_tables.num_kv_cache_groups if block_tables is not None else 0

102 + )

103 + self._global_batch_slot_mappings = (

104 + torch.empty(

105 + num_kv_cache_groups,

106 + max_num_tokens,

107 + dtype=torch.int64,

108 + device=device,

109 + )

110 + if max_num_tokens is not None and num_kv_cache_groups > 0

111 + else None

112 + )

113 + self._gathered_kv_slot_mappings = (

114 + torch.empty(

115 + num_kv_cache_groups,

116 + max_num_tokens * pcp_world_size,

117 + dtype=torch.int64,

118 + device=device,

119 + )

120 + if max_num_tokens is not None and num_kv_cache_groups > 0

121 + else None

122 + )

123 +

124 + @staticmethod

125 + def validate_config(

126 + vllm_config: VllmConfig,

127 + supports_mm_inputs: bool,

128 + ) -> None:

129 + parallel_config = vllm_config.parallel_config

130 + model_config = vllm_config.model_config

131 + pcp_size = parallel_config.prefill_context_parallel_size

132 + if pcp_size <= 1:

133 + return

134 +

135 + if not model_config.use_mla:

136 + raise NotImplementedError("MRV2 PCP currently supports MLA models only.")

137 + if parallel_config.pipeline_parallel_size > 1:

138 + raise NotImplementedError("MRV2 PCP does not support PP yet.")

139 + if model_config.is_encoder_decoder:

140 + raise NotImplementedError(

141 + "MRV2 PCP does not support encoder-decoder models yet."

142 + )

143 + if supports_mm_inputs:

144 + raise NotImplementedError("MRV2 PCP does not support MM inputs yet.")

145 + if vllm_config.lora_config is not None:

146 + raise NotImplementedError("MRV2 PCP does not support LoRA yet.")

147 + if vllm_config.speculative_config is not None:

148 + raise NotImplementedError(

149 + "MRV2 PCP does not support speculative decoding yet."

150 + )

151 + is_sparse_mla = hasattr(model_config.hf_text_config, "index_topk")

152 + if (

153 + is_sparse_mla

154 + and vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

155 + ):

156 + raise NotImplementedError(

157 + "MRV2 sparse MLA PCP does not support CUDA graphs yet. "

158 + "Set -cc.cudagraph_mode=NONE."

159 + )

160 + if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs():

161 + raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.")

162 +

163 + @staticmethod

164 + def _reorder_segments(

165 + segments: list[RankSegment],

166 + num_computed_tokens: np.ndarray,

167 + is_prefilling: np.ndarray,

168 + query_start_loc_np: np.ndarray,

169 + ) -> list[RankSegment]:

170 + """Move pure prefills last to match the batch ordering expected by

171 + attention backends like MLA and sparse MLA.

172 + """

173 +

174 + def is_pure_prefill(segment: RankSegment) -> bool:

175 + req_idx = segment.global_batch_req_idx

176 + start_pos = (

177 + num_computed_tokens[req_idx]

178 + + segment.global_batch_slice.start

179 + - query_start_loc_np[req_idx]

180 + )

181 + return is_prefilling[req_idx] and start_pos == 0

182 +

183 + segments.sort(key=is_pure_prefill)

184 + rank_offset = 0

185 + for index, segment in enumerate(segments):

186 + segments[index] = replace(

187 + segment,

188 + rank_local_batch_slice=slice(

189 + rank_offset, rank_offset + segment.num_tokens

190 + ),

191 + )

192 + rank_offset += segment.num_tokens

193 + return segments

194 +

195 + def _get_rank_segments(

196 + self,

197 + rank: int,

198 + num_scheduled_tokens: np.ndarray,

199 + num_computed_tokens: np.ndarray,

200 + is_prefilling: np.ndarray,

201 + query_start_loc_np: np.ndarray,

202 + ) -> list[RankSegment]:

203 + """Build one rank's attention-compatible DualChunkSwap rows.

204 +

205 + PCP=4 partitions each prefill into eight chunks:

206 +

207 + full: | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |

208 + rank 0: 0 7

209 + rank 1: 1 6

210 + rank 2: 2 5

211 + rank 3: 3 4

212 + """

213 + rank_segments = []

214 + rank_offset = 0

215 + num_chunks = 2 * self.pcp_world_size

216 + for global_batch_req_idx, num_tokens in enumerate(num_scheduled_tokens):

217 + query_len = int(num_tokens)

218 + if query_len == 0:

219 + continue

220 + global_batch_start = int(query_start_loc_np[global_batch_req_idx])

221 + chunk_indices: tuple[int, ...]

222 + if bool(is_prefilling[global_batch_req_idx]):

223 + chunk_size = (query_len + num_chunks - 1) // num_chunks

224 + chunk_indices = (rank, num_chunks - 1 - rank)

225 + else: # decodes are replicated

226 + chunk_size = query_len

227 + chunk_indices = (0,)

228 +

229 + for chunk_idx in chunk_indices:

230 + chunk_offset = chunk_idx * chunk_size

231 + chunk_len = min(chunk_size, query_len - chunk_offset)

232 + if chunk_len <= 0:

233 + continue

234 + chunk_start = global_batch_start + chunk_offset

235 + rank_segments.append(

236 + RankSegment(

237 + global_batch_req_idx=global_batch_req_idx,

238 + global_batch_slice=slice(chunk_start, chunk_start + chunk_len),

239 + rank_local_batch_slice=slice(

240 + rank_offset, rank_offset + chunk_len

241 + ),

242 + )

243 + )

244 + rank_offset += chunk_len

245 + return self._reorder_segments(

246 + rank_segments,

247 + num_computed_tokens,

248 + is_prefilling,

249 + query_start_loc_np,

250 + )

251 +

252 + def _build_batch_layout(

253 + self,

254 + num_scheduled_tokens: np.ndarray,

255 + num_computed_tokens: np.ndarray,

256 + is_prefilling: np.ndarray,

257 + query_start_loc_np: np.ndarray,

258 + ) -> tuple[list[list[RankSegment]], list[int]]:

259 + segments_by_rank = []

260 + per_rank_num_tokens = []

261 + for rank in range(self.pcp_world_size):

262 + segments = self._get_rank_segments(

263 + rank,

264 + num_scheduled_tokens,

265 + num_computed_tokens,

266 + is_prefilling,

267 + query_start_loc_np,

268 + )

269 + num_rank_tokens = sum(segment.num_tokens for segment in segments)

270 + segments_by_rank.append(segments)

271 + per_rank_num_tokens.append(num_rank_tokens)

272 +

273 + # PCP=2 example:

274 + # global batch: [A B C D E F G]

275 + # rank 0 / rank 1: [A B G] / [C D E F]

276 + # padded gathered: [A B G _ | C D E F]

277 + # hidden_restore_idx: [0, 1, 4, 5, 6, 7, 2]

278 + # padded_gather_idx: [0, 1, 6, 0, 2, 3, 4, 5]

279 + # Therefore global = gathered[hidden_restore_idx] and

280 + # padded_gathered = global[padded_gather_idx].

281 + hidden_restore_idx = np.empty(int(query_start_loc_np[-1]), dtype=np.int64)

282 + padded_num_tokens = max(per_rank_num_tokens)

283 + num_expanded_tokens = padded_num_tokens * self.pcp_world_size

284 + padded_gather_idx = np.zeros(num_expanded_tokens, dtype=np.int64)

285 + gathered_kv_write_mask = np.zeros(num_expanded_tokens, dtype=np.bool_)

286 + for rank, segments in enumerate(segments_by_rank):

287 + expanded_rank_offset = rank * padded_num_tokens

288 + for segment in segments:

289 + padded_gathered_slice = slice(

290 + expanded_rank_offset + segment.rank_local_batch_slice.start,

291 + expanded_rank_offset + segment.rank_local_batch_slice.stop,

292 + )

293 + padded_gather_idx[padded_gathered_slice] = np.arange(

294 + segment.global_batch_slice.start,

295 + segment.global_batch_slice.stop,

296 + dtype=np.int64,

297 + )

298 + # Cache insertion pairs one slot entry with each rank's local decode.

299 + if not bool(is_prefilling[segment.global_batch_req_idx]) and rank != 0:

300 + continue

301 + gathered_kv_write_mask[padded_gathered_slice] = True

302 + hidden_restore_idx[segment.global_batch_slice] = np.arange(

303 + padded_gathered_slice.start,

304 + padded_gathered_slice.stop,

305 + dtype=np.int64,

306 + )

307 +

308 + self._hidden_restore_idx = async_copy_to_gpu(

309 + hidden_restore_idx, device=self.device

310 + )

311 + self._padded_gather_idx = async_copy_to_gpu(

312 + padded_gather_idx, device=self.device

313 + )

314 + self._gathered_kv_write_mask = async_copy_to_gpu(

315 + gathered_kv_write_mask, device=self.device

316 + )

317 + return segments_by_rank, per_rank_num_tokens

318 +

319 + def partition_batch(self, input_batch: InputBatch) -> InputBatch:

320 + assert self._req_states is not None

321 + assert self._input_buffers is not None

322 + req_states = self._req_states

323 + input_buffers = self._input_buffers

324 + if input_batch.num_draft_tokens > 0:

325 + raise NotImplementedError("MRV2 PCP does not support spec decode yet.")

326 +

327 + global_batch = input_batch

328 + self._global_batch = global_batch

329 +

330 + num_scheduled_tokens = global_batch.num_scheduled_tokens

331 + num_computed_tokens = global_batch.num_computed_tokens_np

332 + is_prefilling = global_batch.is_prefilling_np

333 +

334 + segments_by_rank, per_rank_num_tokens = self._build_batch_layout(

335 + num_scheduled_tokens,

336 + num_computed_tokens,

337 + is_prefilling,

338 + global_batch.query_start_loc_np,

339 + )

340 +

341 + local_segments = segments_by_rank[self.pcp_rank]

342 + if not local_segments:

343 + local_segments = [

344 + RankSegment(

345 + global_batch_req_idx=0,

346 + global_batch_slice=slice(0, 0),

347 + rank_local_batch_slice=slice(0, 0),

348 + )

349 + ]

350 +

351 + num_local_reqs = len(local_segments)

352 + if num_local_reqs > input_buffers.max_num_reqs:

353 + raise RuntimeError(

354 + "PCP local request count exceeds the MRV2 input buffer size: "

355 + f"{num_local_reqs} > {input_buffers.max_num_reqs}."

356 + )

357 +

358 + local_to_global_batch_req_idx_np = np.fromiter(

359 + (segment.global_batch_req_idx for segment in local_segments),

360 + dtype=np.int32,

361 + count=num_local_reqs,

362 + )

363 + local_start_pos_np = np.fromiter(

364 + (

365 + num_computed_tokens[segment.global_batch_req_idx]

366 + + segment.global_batch_slice.start

367 + - global_batch.query_start_loc_np[segment.global_batch_req_idx]

368 + for segment in local_segments

369 + ),

370 + dtype=np.int32,

371 + count=num_local_reqs,

372 + )

373 + local_num_scheduled_tokens = np.fromiter(

374 + (segment.num_tokens for segment in local_segments),

375 + dtype=np.int32,

376 + count=num_local_reqs,

377 + )

378 + local_to_global_req_idx_np = global_batch.idx_mapping_np[

379 + local_to_global_batch_req_idx_np

380 + ]

381 + local_req_ids = [

382 + global_batch.req_ids[global_batch_req_idx]

383 + for global_batch_req_idx in local_to_global_batch_req_idx_np

384 + ]

385 +

386 + num_local_tokens = int(local_num_scheduled_tokens.sum())

387 + num_local_tokens_padded = max(per_rank_num_tokens)

388 + fresh_prefills = int(

389 + np.count_nonzero(is_prefilling & (num_computed_tokens == 0))

390 + )

391 + continued_prefills = int(

392 + np.count_nonzero(is_prefilling & (num_computed_tokens > 0))

393 + )

394 + logger.debug(

395 + "PCP batch: rank=%d global_batch_reqs=%d fresh_prefills=%d "

396 + "continued_prefills=%d decodes=%d local_reqs=%d "

397 + "local_tokens=%d per_rank_tokens=%s",

398 + self.pcp_rank,

399 + global_batch.num_reqs,

400 + fresh_prefills,

401 + continued_prefills,

402 + global_batch.num_reqs - fresh_prefills - continued_prefills,

403 + num_local_reqs,

404 + num_local_tokens,

405 + per_rank_num_tokens,

406 + )

407 + if num_local_tokens_padded > input_buffers.max_num_tokens:

408 + raise RuntimeError(

409 + "PCP local token count exceeds the MRV2 input buffer size: "

410 + f"{num_local_tokens_padded} > {input_buffers.max_num_tokens}."

411 + )

412 + rank_token_start = self.pcp_rank * num_local_tokens_padded

413 + assert self._padded_gather_idx is not None

414 + local_gather_idx = self._padded_gather_idx[

415 + rank_token_start : rank_token_start + num_local_tokens_padded

416 + ]

417 + torch.index_select(

418 + global_batch.input_ids,

419 + 0,

420 + local_gather_idx,

421 + out=input_buffers.input_ids[:num_local_tokens_padded],

422 + )

423 +

424 + local_query_start_loc_np = np.empty(

425 + input_buffers.max_num_reqs + 1, dtype=np.int32

426 + )

427 + local_query_start_loc_np[0] = 0

428 + local_query_start_loc_out = local_query_start_loc_np[1 : num_local_reqs + 1]

429 + np.cumsum(local_num_scheduled_tokens, out=local_query_start_loc_out)

430 + local_query_start_loc_np[num_local_reqs + 1 :] = num_local_tokens

431 + async_copy_to_gpu(local_query_start_loc_np, out=input_buffers.query_start_loc)

432 + local_query_start_loc = input_buffers.query_start_loc[: num_local_reqs + 1]

433 +

434 + local_to_global_req_idx = async_copy_to_gpu(

435 + local_to_global_req_idx_np, device=self.device

436 + )

437 + local_start_pos = async_copy_to_gpu(local_start_pos_np, device=self.device)

438 +

439 + assert self._local_req_idx is not None

440 + prepare_pos_seq_lens(

441 + self._local_req_idx[:num_local_reqs],

442 + local_query_start_loc,

443 + local_start_pos,

444 + input_buffers.positions,

445 + input_buffers.seq_lens[:num_local_reqs],

446 + )

447 + seq_lens = input_buffers.seq_lens[:num_local_reqs]

448 + is_padding = input_buffers.is_padding[:num_local_tokens_padded]

449 + is_padding[:num_local_tokens].fill_(False)

450 + is_padding[num_local_tokens:].fill_(True)

451 + if num_local_tokens_padded > num_local_tokens:

452 + input_buffers.input_ids[:num_local_tokens_padded].masked_fill_(

453 + is_padding, 0

454 + )

455 + input_buffers.positions[:num_local_tokens_padded].masked_fill_(

456 + is_padding, 0

457 + )

458 +

459 + total_num_logits = num_local_reqs if num_local_tokens > 0 else 0

460 + if total_num_logits > 0:

461 + cu_num_logits_np = np.arange(num_local_reqs + 1, dtype=np.int32)

462 + cu_num_logits = torch.arange(

463 + num_local_reqs + 1, device=self.device, dtype=torch.int32

464 + )

465 + else:

466 + cu_num_logits_np = np.zeros(num_local_reqs + 1, dtype=np.int32)

467 + cu_num_logits = torch.zeros(

468 + num_local_reqs + 1, device=self.device, dtype=torch.int32

469 + )

470 + logits_indices = combine_sampled_and_draft_tokens(

471 + input_buffers.input_ids,

472 + local_to_global_req_idx,

473 + req_states.last_sampled_tokens,

474 + local_query_start_loc,

475 + seq_lens,

476 + req_states.prefill_len.gpu,

477 + req_states.draft_tokens,

478 + cu_num_logits,

479 + total_num_logits,

480 + 1,

481 + )

482 +

483 + local_prefill_len_np = global_batch.prefill_len_np[

484 + local_to_global_batch_req_idx_np

485 + ]

486 + local_num_computed_prefill_tokens_np = np.minimum(

487 + local_start_pos_np, local_prefill_len_np

488 + )

489 + local_is_prefilling_np = (

490 + local_num_computed_prefill_tokens_np < local_prefill_len_np

491 + )

492 + seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32)

493 + seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens

494 +

495 + dcp_local_seq_lens = None

496 + if self.dcp_world_size > 1:

497 + prepare_dcp_local_seq_lens(

498 + input_buffers.dcp_local_seq_lens,

499 + seq_lens,

500 + num_local_reqs,

501 + self.dcp_world_size,

502 + self.dcp_rank,

503 + self.cp_interleave,

504 + )

505 + dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_local_reqs]

506 +

507 + return replace(

508 + input_batch,

509 + req_ids=local_req_ids,

510 + num_reqs=num_local_reqs,

511 + num_reqs_after_padding=num_local_reqs,

512 + idx_mapping=local_to_global_req_idx,

513 + idx_mapping_np=local_to_global_req_idx_np,

514 + expanded_idx_mapping=local_to_global_req_idx,

515 + expanded_local_pos=torch.zeros(

516 + num_local_reqs, dtype=torch.int32, device=self.device

517 + ),

518 + num_scheduled_tokens=local_num_scheduled_tokens,

519 + num_tokens=num_local_tokens,

520 + num_tokens_after_padding=num_local_tokens_padded,

521 + num_draft_tokens=0,

522 + num_draft_tokens_per_req=None,

523 + query_start_loc=local_query_start_loc,

524 + query_start_loc_np=local_query_start_loc_np[: num_local_reqs + 1],

525 + seq_lens=seq_lens,

526 + seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np),

527 + dcp_local_seq_lens=dcp_local_seq_lens,

528 + num_computed_tokens_np=local_start_pos_np,

529 + prefill_len_np=local_prefill_len_np,

530 + num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np,

531 + is_prefilling_np=local_is_prefilling_np,

532 + max_seq_len_np=global_batch.max_seq_len_np[local_to_global_batch_req_idx_np]

533 + if global_batch.max_seq_len_np is not None

534 + else None,

535 + input_ids=input_buffers.input_ids[:num_local_tokens_padded],

536 + positions=input_buffers.positions[:num_local_tokens_padded],

537 + is_padding=is_padding,

538 + logits_indices=logits_indices,

539 + cu_num_logits=cu_num_logits,

540 + cu_num_logits_np=cu_num_logits_np,

541 + prompt_lens=None,

542 + )

543 +

544 + def prepare_attn(

545 + self, input_batch: InputBatch

546 + ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:

547 + assert self._block_tables is not None

548 + assert self._local_block_tables is not None

549 + assert self._local_block_table_ptrs is not None

550 + block_tables = self._block_tables.gather_block_tables(

551 + input_batch.idx_mapping,

552 + input_batch.num_reqs_after_padding,

553 + out=self._local_block_tables,

554 + out_ptrs=self._local_block_table_ptrs,

555 + )

556 + slot_mappings = self.prepare_slot_mappings()

557 + return block_tables, slot_mappings

558 +

559 + def prepare_slot_mappings(self) -> torch.Tensor:

560 + assert self._block_tables is not None

561 + assert self._global_batch_slot_mappings is not None

562 + assert self._global_batch is not None

563 + global_batch = self._global_batch

564 + global_batch_slot_mappings = self._block_tables.compute_slot_mappings(

565 + global_batch.idx_mapping,

566 + global_batch.query_start_loc,

567 + global_batch.positions,

568 + global_batch.num_tokens,

569 + out=self._global_batch_slot_mappings,

570 + )

571 + return self._convert_to_gathered_slot_mappings(global_batch_slot_mappings)

572 +

573 + def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:

574 + assert self._gathered_kv_slot_mappings is not None

575 + self._gathered_kv_slot_mappings.fill_(PAD_SLOT_ID)

576 + return self._gathered_kv_slot_mappings[:, : num_tokens * self.pcp_world_size]

577 +

578 + def _convert_to_gathered_slot_mappings(

579 + self,

580 + global_batch_slot_mappings: torch.Tensor,

581 + ) -> torch.Tensor:

582 + assert self._padded_gather_idx is not None

583 + assert self._gathered_kv_write_mask is not None

584 + padded_gather_idx = self._padded_gather_idx

585 + num_expanded_tokens = padded_gather_idx.shape[0]

586 + if self._gathered_kv_slot_mappings is None:

587 + self._gathered_kv_slot_mappings = global_batch_slot_mappings.new_empty(

588 + global_batch_slot_mappings.shape[0], num_expanded_tokens

589 + )

590 + gathered_kv_slot_mappings = self._gathered_kv_slot_mappings[

591 + :, :num_expanded_tokens

592 + ]

593 + torch.index_select(

594 + global_batch_slot_mappings,

595 + 1,

596 + padded_gather_idx,

597 + out=gathered_kv_slot_mappings,

598 + )

599 + torch.where(

600 + self._gathered_kv_write_mask.unsqueeze(0),

601 + gathered_kv_slot_mappings,

602 + self._pad_slot_id,

603 + out=gathered_kv_slot_mappings,

604 + )

605 + return gathered_kv_slot_mappings

606 +

607 + def restore_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor:

608 + if self._hidden_restore_idx is None:

609 + return hidden_states

610 + gathered = get_pcp_group().all_gather(hidden_states, dim=0)

611 + return gathered[self._hidden_restore_idx]

612 +

613 + def restore_for_sampling(

614 + self,

615 + hidden_states: torch.Tensor,

616 + ) -> tuple[torch.Tensor, InputBatch]:

617 + assert self._global_batch is not None

618 + return self.restore_hidden_states(hidden_states), self._global_batch

619 +

620 +

621 +def maybe_partition_pcp_batch(

622 + manager: PCPManager | None,

623 + input_batch: InputBatch,

624 +) -> InputBatch:

625 + if manager is None:

626 + return input_batch

627 + return manager.partition_batch(input_batch)

628 +

629 +

630 +def maybe_get_pcp_dummy_slot_mappings(

631 + manager: PCPManager | None,

632 + block_tables: BlockTables,

633 + num_tokens: int,

634 +) -> torch.Tensor:

635 + if manager is None:

636 + return block_tables.get_dummy_slot_mappings(num_tokens)

637 + return manager.get_dummy_slot_mappings(num_tokens)

638 +

639 +

640 +def maybe_restore_pcp_for_sampling(

641 + manager: PCPManager | None,

642 + hidden_states: torch.Tensor | None,

643 + input_batch: InputBatch,

644 +) -> tuple[torch.Tensor, InputBatch]:

645 + assert hidden_states is not None

646 + if manager is None:

647 + return hidden_states, input_batch

648 + return manager.restore_for_sampling(hidden_states)

649 +

650 +

651 +def maybe_build_pcp_manager(

652 + vllm_config: VllmConfig,

653 + device: torch.device,

654 + supports_mm_inputs: bool,

655 + req_states: RequestState,

656 + block_tables: BlockTables,

657 +) -> PCPManager | None:

658 + parallel_config = vllm_config.parallel_config

659 + pcp_size = parallel_config.prefill_context_parallel_size

660 + if pcp_size <= 1:

661 + return None

662 +

663 + PCPManager.validate_config(vllm_config, supports_mm_inputs)

664 +

665 + pcp_rank = get_pcp_group().rank_in_group

666 + dcp_size = parallel_config.decode_context_parallel_size

667 + dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0

668 +

669 + return PCPManager(

670 + pcp_world_size=pcp_size,

671 + pcp_rank=pcp_rank,

672 + device=device,

673 + req_states=req_states,

674 + max_num_reqs=vllm_config.scheduler_config.max_num_seqs,

675 + max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens,

676 + block_tables=block_tables,

677 + dcp_world_size=dcp_size,

678 + dcp_rank=dcp_rank,

679 + cp_interleave=parallel_config.cp_kv_cache_interleave_size,

680 + )