@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8486f1ebfe

[Perf][PCP] Shard decode requests across PCP ranks (#52162)

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

Qiu Chunshuo committed Oct 2, 2026 at 09:20 UTC 8486f1ebfe6e39a5513a8636e7f07ed9746f8be5
13 files changed +406 -33
tests/model_executor/layers/test_mla_short_prefill_indexer.py
+2
@@ -205,6 +205,7 @@ def test_short_prefill_updates_k_cache_before_scoring_decision(
205 topk_indices,
206 False,
207 False,
208 + False,
209 MLA_LAYER,
210 )
211
@@ -259,6 +260,7 @@ def test_skipped_k_cache_insert_accepts_no_k(
260 topk_indices,
261 True,
262 False,
263 + False,
264 "",
265 )
266
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+73
@@ -15,6 +15,7 @@ from vllm.models.deepseek_v41.sparse_mla import (
15 DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend,
16 )
17 from vllm.v1.attention.backend import CommonAttentionMetadata
18 +from vllm.v1.attention.backends.mla import indexer as indexer_module
19 from vllm.v1.attention.backends.mla.compressor_utils import (
20 CompressedSlotMappingKernel,
21 get_compressed_slot_mapping,
@@ -295,6 +296,78 @@ def test_index_conversion_warmup_uses_physical_block_stride():
296 assert {key.block_stride_rows for key in keys} == {4096}
297
298
299 +def test_zero_token_pcp_rank_participates_in_compressed_mapping_gather(monkeypatch):
300 + builder = object.__new__(DeepseekV32IndexerMetadataBuilder)
301 + builder.compress_ratio = 4
302 + builder.use_pcp = True
303 + builder.pcp_world_size = 4
304 + builder.pcp_rank = 0
305 + builder.kernel_block_size = None
306 + builder.kv_cache_spec = SimpleNamespace(block_size=64, num_states=64)
307 + builder.compressed_slot_mapping_buffer = torch.zeros(8, dtype=torch.int64)
308 +
309 + def fake_get_compressed_slot_mapping(
310 + num_tokens,
311 + slot_mapping,
312 + query_start_loc,
313 + seq_lens,
314 + block_table,
315 + block_size,
316 + compress_ratio,
317 + out,
318 + ):
319 + assert num_tokens == 0
320 + assert block_size == 64
321 + assert compress_ratio == 4
322 + out.fill_(-1)
323 + return out[:num_tokens]
324 +
325 + gathered_slot_mapping = torch.tensor([123, -1, -1, -1], dtype=torch.int64)
326 +
327 + class FakePCPGroup:
328 + def __init__(self):
329 + self.calls = 0
330 +
331 + def all_gather(self, tensor, dim=0):
332 + self.calls += 1
333 + assert dim == 0
334 + torch.testing.assert_close(tensor, torch.tensor([-1]))
335 + return gathered_slot_mapping
336 +
337 + fake_pcp_group = FakePCPGroup()
338 + monkeypatch.setattr(
339 + indexer_module,
340 + "get_compressed_slot_mapping",
341 + fake_get_compressed_slot_mapping,
342 + )
343 + monkeypatch.setattr(indexer_module, "get_pcp_group", lambda: fake_pcp_group)
344 +
345 + query_start_loc = torch.tensor([0, 0], dtype=torch.int32)
346 + seq_lens = torch.tensor([16], dtype=torch.int32)
347 + common = CommonAttentionMetadata(
348 + query_start_loc=query_start_loc,
349 + query_start_loc_cpu=query_start_loc,
350 + seq_lens=seq_lens,
351 + seq_lens_cpu_upper_bound=seq_lens,
352 + num_reqs=1,
353 + num_actual_tokens=0,
354 + max_query_len=0,
355 + max_seq_len=16,
356 + block_table_tensor=torch.tensor([[0]], dtype=torch.int32),
357 + slot_mapping=torch.tensor([11, 22, 33, 44], dtype=torch.int64),
358 + causal=True,
359 + )
360 +
361 + metadata = builder.build(common_prefix_len=0, common_attn_metadata=common)
362 +
363 + assert fake_pcp_group.calls == 1
364 + assert metadata.num_decodes == 0
365 + assert metadata.num_decode_tokens == 0
366 + assert metadata.num_prefills == 0
367 + assert metadata.num_prefill_tokens == 0
368 + torch.testing.assert_close(metadata.slot_mapping, gathered_slot_mapping)
369 +
370 +
371 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
372 def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_num_states():
373 """Regression test: DeepseekV4 compression path must compute slot_mapping from
tests/v1/attention/test_sparse_mla_backends.py
+1
@@ -2499,6 +2499,7 @@ def test_hisparse_kv_update_writes_resident_and_staging_caches():
2499 layer = SimpleNamespace(
2500 hisparse_cache=cache_handle,
2501 use_pcp=False,
2502 + pcp_shard_decode_requests=False,
2503 impl=impl,
2504 )
2505 mla_attention.MLAAttention.update_kv_cache(
tests/v1/cudagraph/test_cudagraph_manager.py
+1
@@ -132,6 +132,7 @@ def test_piecewise_capture_uses_pcp_dummy_slot_mappings():
132 pcp_world_size=pcp_world_size,
133 pcp_rank=0,
134 device=torch.device("cpu"),
135 + shard_decode_requests=False,
136 max_num_reqs=num_reqs,
137 max_num_tokens=num_tokens,
138 block_tables=pcp_block_tables,
tests/v1/worker/test_gpu_pcp_manager.py
+220 -10
@@ -7,8 +7,9 @@ import numpy as np
7 import pytest
8 import torch
9
10 -from vllm.config import CUDAGraphMode
11 -from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
10 +from vllm.config import CUDAGraphMode, ParallelConfig
11 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens
12 +from vllm.v1.attention.ops import pcp as attention_pcp
13 from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
14 from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module
15 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
@@ -53,6 +54,7 @@ def _make_capture_manager(block_table: torch.Tensor):
54 pcp_world_size=2,
55 pcp_rank=0,
56 device=torch.device("cpu"),
57 + shard_decode_requests=False,
58 max_num_reqs=4,
59 max_num_tokens=8,
60 block_tables=block_tables,
@@ -93,11 +95,12 @@ def test_validate_config_dcp_for_dspark(draft_uses_mla, dcp_size):
95 PCPManager.validate_config(config, supports_mm_inputs=False)
96
97
96 -def test_replicated_decode_piecewise_graph_padding(monkeypatch):
98 +def test_sharded_decode_piecewise_graph_padding(monkeypatch):
99 manager = PCPManager(
100 pcp_world_size=2,
101 pcp_rank=0,
102 device=torch.device("cpu"),
103 + shard_decode_requests=True,
104 dcp_world_size=1,
105 )
106 monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu)
@@ -110,19 +113,19 @@ def test_replicated_decode_piecewise_graph_padding(monkeypatch):
113 padded_num_tokens=4,
114 )
115
113 - assert per_rank_num_tokens == [3, 3]
116 + assert per_rank_num_tokens == [2, 1]
117 request_indices = [
118 [segment.global_batch_req_idx for segment in rank] for rank in segments_by_rank
119 ]
117 - assert request_indices == [[0, 1, 2], [0, 1, 2]]
118 - assert torch.equal(manager._hidden_restore_idx, torch.tensor([0, 1, 2]))
120 + assert request_indices == [[0, 2], [1]]
121 + assert torch.equal(manager._hidden_restore_idx, torch.tensor([0, 4, 1]))
122 assert torch.equal(
123 manager._padded_gather_idx,
121 - torch.tensor([0, 1, 2, 0, 0, 1, 2, 0]),
124 + torch.tensor([0, 2, 0, 0, 1, 0, 0, 0]),
125 )
126 assert torch.equal(
127 manager._gathered_kv_write_mask,
125 - torch.tensor([True, True, True, False, False, False, False, False]),
128 + torch.tensor([True, True, False, False, True, False, False, False]),
129 )
130
131
@@ -131,6 +134,7 @@ def test_input_buffers_are_exposed_for_cudagraph_capture():
134 pcp_world_size=2,
135 pcp_rank=0,
136 device=torch.device("cpu"),
137 + shard_decode_requests=False,
138 max_num_reqs=4,
139 max_num_tokens=8,
140 )
@@ -148,7 +152,7 @@ def test_input_buffers_are_exposed_for_cudagraph_capture():
152 (2, [7], [True], 4),
153 (2, [3], [False], 3),
154 (2, [3, 8], [False, True], 7),
151 - (4, [2, 9], [False, True], 5),
155 + (4, [2, 9], [False, True], 4),
156 ],
157 )
158 def test_num_tokens_for_dispatch_uses_largest_pcp_rank(
@@ -158,6 +162,7 @@ def test_num_tokens_for_dispatch_uses_largest_pcp_rank(
162 pcp_world_size=pcp_world_size,
163 pcp_rank=0,
164 device=torch.device("cpu"),
165 + shard_decode_requests=True,
166 )
167
168 actual = manager.get_num_tokens_for_dispatch(
@@ -173,6 +178,7 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
178 pcp_world_size=2,
179 pcp_rank=0,
180 device=torch.device("cpu"),
181 + shard_decode_requests=True,
182 dcp_world_size=1,
183 )
184 monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu)
@@ -183,7 +189,7 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
189 num_computed_tokens=np.full(3, 16, dtype=np.int32),
190 is_prefilling=np.zeros(3, dtype=np.bool_),
191 query_start_loc_np=np.arange(4, dtype=np.int32),
186 - padded_num_tokens=2,
192 + padded_num_tokens=1,
193 )
194
195
@@ -207,6 +213,7 @@ def test_partition_padding_is_derived_from_batch_descriptor(
213 pcp_world_size=2,
214 pcp_rank=0,
215 device=device,
216 + shard_decode_requests=False,
217 max_num_reqs=4,
218 max_num_tokens=8,
219 )
@@ -288,6 +295,189 @@ def test_dummy_context_updates_pcp_local_block_tables():
295 assert torch.all(global_block_table == -1)
296
297
298 +def _rank_request_ids(
299 + manager: PCPManager,
300 + rank: int,
301 + req_ids: list[str],
302 + *,
303 + is_prefilling: np.ndarray | None = None,
304 +) -> list[str]:
305 + num_reqs = len(req_ids)
306 + if is_prefilling is None:
307 + is_prefilling = np.zeros(num_reqs, dtype=np.bool_)
308 + segments = manager._get_rank_segments(
309 + rank=rank,
310 + num_scheduled_tokens=np.ones(num_reqs, dtype=np.int32),
311 + is_prefilling=is_prefilling,
312 + query_start_loc_np=np.arange(num_reqs + 1, dtype=np.int32),
313 + )
314 + return [req_ids[segment.global_batch_req_idx] for segment in segments]
315 +
316 +
317 +def test_pcp_only_decode_requests_are_round_robin_balanced_each_step():
318 + manager = PCPManager(
319 + pcp_world_size=4,
320 + pcp_rank=0,
321 + device=torch.device("cpu"),
322 + shard_decode_requests=True,
323 + dcp_world_size=1,
324 + )
325 + req_ids = [f"request-{idx}" for idx in range(18)]
326 +
327 + owners: dict[str, int] = {}
328 + for rank in range(manager.pcp_world_size):
329 + for req_id in _rank_request_ids(manager, rank, req_ids):
330 + assert req_id not in owners
331 + owners[req_id] = rank
332 +
333 + assert owners == {
334 + req_id: index % manager.pcp_world_size for index, req_id in enumerate(req_ids)
335 + }
336 + counts = [list(owners.values()).count(rank) for rank in range(4)]
337 + assert max(counts) - min(counts) == 1
338 +
339 + reordered_req_ids = req_ids[::2] + req_ids[1::2]
340 + reordered_owners = {
341 + req_id: rank
342 + for rank in range(manager.pcp_world_size)
343 + for req_id in _rank_request_ids(manager, rank, reordered_req_ids)
344 + }
345 + assert reordered_owners == {
346 + req_id: index % manager.pcp_world_size
347 + for index, req_id in enumerate(reordered_req_ids)
348 + }
349 + assert reordered_owners != owners
350 +
351 +
352 +def test_decode_requests_remain_replicated_when_dcp_is_enabled():
353 + manager = PCPManager(
354 + pcp_world_size=2,
355 + pcp_rank=0,
356 + device=torch.device("cpu"),
357 + shard_decode_requests=False,
358 + dcp_world_size=2,
359 + )
360 + req_ids = ["request-a", "request-b", "request-c"]
361 +
362 + assert _rank_request_ids(manager, 0, req_ids) == req_ids
363 + assert _rank_request_ids(manager, 1, req_ids) == req_ids
364 +
365 +
366 +def test_decode_sharding_allows_ranks_with_no_owned_request():
367 + manager = PCPManager(
368 + pcp_world_size=4,
369 + pcp_rank=0,
370 + device=torch.device("cpu"),
371 + shard_decode_requests=True,
372 + dcp_world_size=1,
373 + )
374 + req_ids = ["request-a"]
375 +
376 + assert _rank_request_ids(manager, 0, req_ids) == req_ids
377 + assert _rank_request_ids(manager, 1, req_ids) == []
378 + assert _rank_request_ids(manager, 2, req_ids) == []
379 + assert _rank_request_ids(manager, 3, req_ids) == []
380 +
381 +
382 +def test_prefill_partitioning_is_preserved_with_sharded_decode():
383 + manager = PCPManager(
384 + pcp_world_size=2,
385 + pcp_rank=0,
386 + device=torch.device("cpu"),
387 + shard_decode_requests=True,
388 + dcp_world_size=1,
389 + )
390 + segments_by_rank = [
391 + manager._get_rank_segments(
392 + rank=rank,
393 + num_scheduled_tokens=np.array([8, 1, 0, 1], dtype=np.int32),
394 + is_prefilling=np.array([True, False, False, False]),
395 + query_start_loc_np=np.array([0, 8, 9, 9, 10], dtype=np.int32),
396 + )
397 + for rank in range(manager.pcp_world_size)
398 + ]
399 +
400 + prefill_tokens = sorted(
401 + token_idx
402 + for segments in segments_by_rank
403 + for segment in segments
404 + if segment.global_batch_req_idx == 0
405 + for token_idx in range(
406 + segment.global_batch_slice.start, segment.global_batch_slice.stop
407 + )
408 + )
409 + decode_owners = {
410 + segment.global_batch_req_idx: rank
411 + for rank, segments in enumerate(segments_by_rank)
412 + for segment in segments
413 + if segment.global_batch_req_idx in (1, 3)
414 + }
415 +
416 + assert prefill_tokens == list(range(8))
417 + assert decode_owners == {1: 0, 3: 1}
418 +
419 +
420 +def test_sharded_decode_layout_selects_owner_kv_for_replication(monkeypatch):
421 + manager = PCPManager(
422 + pcp_world_size=2,
423 + pcp_rank=0,
424 + device=torch.device("cpu"),
425 + shard_decode_requests=True,
426 + dcp_world_size=1,
427 + )
428 + monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu)
429 + manager._build_batch_layout(
430 + num_scheduled_tokens=np.array([1, 1, 1], dtype=np.int32),
431 + num_computed_tokens=np.array([16, 16, 16], dtype=np.int32),
432 + is_prefilling=np.array([False, False, False]),
433 + query_start_loc_np=np.array([0, 1, 2, 3], dtype=np.int32),
434 + )
435 +
436 + gathered_slot_mapping = manager._convert_to_gathered_slot_mappings(
437 + torch.tensor([[123, 456, 789]], dtype=torch.int64)
438 + )
439 + assert torch.equal(
440 + gathered_slot_mapping,
441 + torch.tensor([[123, 789, 456, PAD_SLOT_ID]], dtype=torch.int64),
442 + )
443 + assert torch.equal(manager._hidden_restore_idx, torch.tensor([0, 2, 1]))
444 +
445 + class FakePCPGroup:
446 + world_size = 2
447 +
448 + def all_gather(self, tensor: torch.Tensor, dim: int = 0) -> torch.Tensor:
449 + assert dim == 0
450 + assert tensor.shape == (2, 1)
451 + return torch.tensor([[11.0], [33.0], [22.0], [0.0]])
452 +
453 + monkeypatch.setattr(attention_pcp, "get_pcp_group", FakePCPGroup)
454 + (gathered_kv,), cache_slot_mapping = attention_pcp._gather_prefill_cache_inputs(
455 + (torch.tensor([[11.0], [33.0]]),),
456 + gathered_slot_mapping[0],
457 + num_decode_tokens=0,
458 + shard_decode_requests=True,
459 + )
460 +
461 + assert torch.equal(gathered_kv, torch.tensor([[11.0], [33.0], [22.0], [0.0]]))
462 + assert torch.equal(cache_slot_mapping, torch.tensor([123, 789, 456, PAD_SLOT_ID]))
463 +
464 +
465 +def test_replicated_draft_cache_inputs_keep_unexpanded_slot_mapping(monkeypatch):
466 + def unexpected_gather():
467 + pytest.fail("replicated draft decode must not gather PCP cache inputs")
468 +
469 + monkeypatch.setattr(attention_pcp, "get_pcp_group", unexpected_gather)
470 + kv = torch.tensor([[11.0], [22.0], [0.0]])
471 + slot_mapping = torch.tensor([123, 456])
472 +
473 + (cache_kv,), cache_slot_mapping = attention_pcp._gather_prefill_cache_inputs(
474 + (kv,), slot_mapping, num_decode_tokens=2, shard_decode_requests=True
475 + )
476 +
477 + assert torch.equal(cache_kv, kv[:2])
478 + assert torch.equal(cache_slot_mapping, slot_mapping)
479 +
480 +
481 def _rank_rows(
482 pcp_rank: int,
483 pcp_world_size: int,
@@ -300,6 +490,7 @@ def _rank_rows(
490 pcp_world_size=pcp_world_size,
491 pcp_rank=pcp_rank,
492 device=torch.device("cpu"),
493 + shard_decode_requests=False,
494 dcp_world_size=2,
495 )
496 query_start_loc_np = np.concatenate([[0], np.cumsum(num_scheduled_tokens)]).astype(
@@ -393,6 +584,7 @@ def test_dcp_replicates_prefills_too_short_to_split(query_len):
584 pcp_world_size=pcp_world_size,
585 pcp_rank=rank,
586 device=torch.device("cpu"),
587 + shard_decode_requests=False,
588 dcp_world_size=2,
589 )
590 assert manager.replicated_requests(num_scheduled_tokens, is_prefilling)[0]
@@ -418,6 +610,7 @@ def test_pcp_first_chunk_row_is_never_short(pcp_world_size, query_len):
610 pcp_world_size=pcp_world_size,
611 pcp_rank=rank,
612 device=torch.device("cpu"),
613 + shard_decode_requests=False,
614 dcp_world_size=pcp_world_size,
615 )
616 chunk_lens = [
@@ -493,6 +686,7 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
686 pcp_world_size=2,
687 pcp_rank=0,
688 device=device,
689 + shard_decode_requests=False,
690 max_num_reqs=4,
691 max_num_tokens=8,
692 dcp_world_size=2,
@@ -541,3 +735,19 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
735 )
736 assert local_batch.dcp_local_seq_lens is not None
737 assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected)
738 +
739 +
740 +@pytest.mark.parametrize(
741 + ("pcp_world_size", "dcp_world_size", "expected"),
742 + [(1, 1, False), (2, 1, True), (2, 2, False)],
743 +)
744 +def test_parallel_config_manages_decode_sharding(
745 + pcp_world_size: int, dcp_world_size: int, expected: bool
746 +):
747 + parallel_config = ParallelConfig(
748 + prefill_context_parallel_size=pcp_world_size,
749 + decode_context_parallel_size=dcp_world_size,
750 + distributed_executor_backend="ray",
751 + )
752 +
753 + assert parallel_config.pcp_shard_decode_requests is expected
vllm/config/parallel.py
+13
@@ -599,6 +599,19 @@ class ParallelConfig:
599 """Process world size across TP, PCP, PP, and DP."""
600 return self.world_size * self.data_parallel_size
601
602 + @property
603 + def pcp_shard_decode_requests(self) -> bool:
604 + """Whether PCP can shard decode requests across its ranks.
605 +
606 + PCP-only execution replicates the KV cache, so each decode request can
607 + have a single PCP owner. DCP shards the KV cache and therefore requires
608 + every decode request to run on every participating DCP rank.
609 + """
610 + return (
611 + self.prefill_context_parallel_size > 1
612 + and self.decode_context_parallel_size == 1
613 + )
614 +
615 @property
616 def use_ubatching(self) -> bool:
617 return self.enable_dbo or self.ubatch_size > 1
vllm/model_executor/layers/attention/mla_attention.py
+2
@@ -609,6 +609,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
609 vllm_config = get_current_vllm_config()
610 parallel_config = vllm_config.parallel_config
611 self.use_pcp = parallel_config.prefill_context_parallel_size > 1
612 + self.pcp_shard_decode_requests = parallel_config.pcp_shard_decode_requests
613 compilation_config = vllm_config.compilation_config
614 if prefix in compilation_config.static_forward_context:
615 raise ValueError(f"Duplicate layer name: {prefix}")
@@ -784,6 +785,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
785 slot_mapping,
786 attn_metadata.num_decode_tokens if attn_metadata is not None else None,
787 self.use_pcp,
788 + pcp_shard_decode_requests=self.pcp_shard_decode_requests,
789 )
790 assert slot_mapping is not None
791 if cache is not None:
vllm/model_executor/layers/sparse_attn_indexer.py
+6
@@ -335,6 +335,7 @@ def sparse_attn_indexer(
335 topk_indices_buffer: torch.Tensor,
336 skip_k_cache_insert: bool,
337 use_pcp: bool,
338 + pcp_shard_decode_requests: bool,
339 dense_mha_metadata_layer_name: LayerNameType,
340 use_fp4_cache: bool = False,
341 dcp_rank: int = 0,
@@ -404,6 +405,7 @@ def sparse_attn_indexer(
405 topk_indices_buffer,
406 skip_k_cache_insert,
407 use_pcp,
408 + pcp_shard_decode_requests,
409 dense_mha_metadata_layer_name,
410 use_fp4_cache,
411 candidate_blocks=candidate_blocks,
@@ -442,6 +444,7 @@ def sparse_attn_indexer(
444 slot_mapping,
445 num_decode_tokens,
446 use_pcp,
447 + pcp_shard_decode_requests=pcp_shard_decode_requests,
448 )
449 # scale_fmt can be None, but the function expects str
450 assert scale_fmt is not None
@@ -817,6 +820,7 @@ def sparse_attn_indexer_fake(
820 topk_indices_buffer: torch.Tensor | None,
821 skip_k_cache_insert: bool,
822 use_pcp: bool,
823 + pcp_shard_decode_requests: bool,
824 dense_mha_metadata_layer_name: LayerNameType,
825 use_fp4_cache: bool = False,
826 dcp_rank: int = 0,
@@ -898,6 +902,7 @@ class SparseAttnIndexer(CustomOp):
902 self.dcp_world_size = parallel_config.decode_context_parallel_size
903 self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
904 self.use_pcp = parallel_config.prefill_context_parallel_size > 1
905 + self.pcp_shard_decode_requests = parallel_config.pcp_shard_decode_requests
906 self._cp_kv_cache_interleave_size: int | None = None
907 if current_platform.is_cuda() and not has_deep_gemm():
908 raise RuntimeError(
@@ -994,6 +999,7 @@ class SparseAttnIndexer(CustomOp):
999 self.topk_indices_buffer,
1000 self.skip_k_cache_insert,
1001 self.use_pcp,
1002 + self.pcp_shard_decode_requests,
1003 _encode_layer_name(self.dense_mha_metadata_layer_name),
1004 self.use_fp4_cache,
1005 self.dcp_rank,
vllm/models/deepseek_v32/attention.py
+2
@@ -500,6 +500,7 @@ class DeepseekV32Attention(MLAAttention):
500 self.topk_indices_buffer,
501 skip_k_cache_insert=not self.use_pcp,
502 use_pcp=self.use_pcp,
503 + pcp_shard_decode_requests=self.pcp_shard_decode_requests,
504 dense_mha_metadata_layer_name=self._dense_mha_metadata_layer_name,
505 dcp_rank=(
506 self.dcp_manager.group.rank_in_group
@@ -535,6 +536,7 @@ class DeepseekV32Attention(MLAAttention):
536 layer_slot_mapping,
537 attn_metadata.num_decode_tokens,
538 True,
539 + pcp_shard_decode_requests=self.pcp_shard_decode_requests,
540 )
541 )
542 self.impl.do_kv_cache_update( # type: ignore[attr-defined]
vllm/v1/attention/backends/mla/flashmla_sparse.py
+8 -1
@@ -489,6 +489,14 @@ class FlashMLASparseMetadataBuilder(
489 num_tokens - metadata.num_decode_tokens,
490 )
491
492 + FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode
493 + # PCP decode sharding can leave a rank with only its collective-padding
494 + # row when the global decode batch is smaller than the PCP world size.
495 + # The row has no actual query tokens and must not be interpreted as a
496 + # zero-length decode request by the sparse FP8 metadata builder.
497 + if num_tokens == 0:
498 + return FP8Meta()
499 +
500 decode_query_len = 0
501 active_num_decodes = num_decodes
502 if num_decodes > 0:
@@ -498,7 +506,6 @@ class FlashMLASparseMetadataBuilder(
506 active_num_decodes = num_decode_tokens // decode_query_len
507 assert active_num_decodes * decode_query_len == num_decode_tokens
508
501 - FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode
509 fp8_metadata = FP8Meta(
510 num_decodes=active_num_decodes,
511 num_prefills=num_prefills,
vllm/v1/attention/backends/mla/indexer.py
+28 -11
@@ -1298,17 +1298,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1298 slot_mapping = common_attn_metadata.slot_mapping
1299 block_table = common_attn_metadata.block_table_tensor
1300 dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens
1301 - num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
1302 - split_decodes_and_prefills(
1303 - common_attn_metadata,
1304 - decode_threshold=self.decode_threshold,
1305 - require_uniform=not (self.use_flattening or self.supports_varlen),
1306 - treat_short_extends_as_decodes=not self.use_pcp,
1307 - )
1308 - )
1309 -
1310 - assert num_decodes + num_prefills == num_reqs
1311 - assert num_decode_tokens + num_prefill_tokens == num_tokens
1301
1302 compressed_slot_mapping = slot_mapping
1303 indexer_block_table = block_table
@@ -1347,6 +1336,34 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1336 dim=0,
1337 )
1338
1339 + # PCP decode sharding keeps a zero-token placeholder on ranks that own
1340 + # no request in a step so collectives retain a uniform rank shape. Do
1341 + # not turn that placeholder into a zero-length indexer decode request.
1342 + if num_tokens == 0:
1343 + return DeepseekV32IndexerMetadata(
1344 + seq_lens=seq_lens,
1345 + max_seq_len=common_attn_metadata.max_seq_len,
1346 + slot_mapping=compressed_slot_mapping,
1347 + num_decodes=0,
1348 + num_decode_tokens=0,
1349 + num_prefills=0,
1350 + num_prefill_tokens=0,
1351 + prefill=None,
1352 + decode=None,
1353 + )
1354 +
1355 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
1356 + split_decodes_and_prefills(
1357 + common_attn_metadata,
1358 + decode_threshold=self.decode_threshold,
1359 + require_uniform=not (self.use_flattening or self.supports_varlen),
1360 + treat_short_extends_as_decodes=not self.use_pcp,
1361 + )
1362 + )
1363 +
1364 + assert num_decodes + num_prefills == num_reqs
1365 + assert num_decode_tokens + num_prefill_tokens == num_tokens
1366 +
1367 prefill_metadata = None
1368 if num_prefills > 0:
1369 compressed_seq_lens = (
vllm/v1/attention/ops/pcp.py
+24 -8
@@ -12,29 +12,42 @@ def _gather_prefill_cache_inputs(
12 tensors: tuple[torch.Tensor, ...],
13 slot_mapping: torch.Tensor,
14 num_decode_tokens: int,
15 + shard_decode_requests: bool = False,
16 ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
16 - """Keep replicated decode writes local and gather partitioned prefills."""
17 + """Gather PCP cache inputs while preserving replicated KV-cache state.
18 +
19 + PCP-only execution shards decode requests across ranks. In that mode each
20 + rank must gather the other owners' decode KV as well as partitioned prefill
21 + KV so every PCP rank retains a complete cache replica. DCP execution keeps
22 + decode requests replicated and uses the legacy prefill-only gather.
23 + """
24 local_num_tokens = tensors[0].shape[0]
25 assert all(tensor.shape[0] == local_num_tokens for tensor in tensors)
26 assert 0 <= num_decode_tokens <= local_num_tokens
27
28 # Replicated draft decodes use unexpanded slot mappings, even with DP padding.
22 - if (
23 - num_decode_tokens == local_num_tokens
24 - or slot_mapping.shape[0] <= local_num_tokens
25 - ):
29 + if slot_mapping.shape[0] <= local_num_tokens:
30 return (
31 tuple(tensor[:num_decode_tokens] for tensor in tensors),
32 slot_mapping[:num_decode_tokens],
33 )
34
35 pcp_group = get_pcp_group()
36 + pcp_size = pcp_group.world_size
37 + gathered_slot_mapping = slot_mapping[: pcp_size * local_num_tokens]
38 + if shard_decode_requests:
39 + gathered_inputs = tuple(
40 + pcp_group.all_gather(tensor.contiguous(), dim=0) for tensor in tensors
41 + )
42 + return gathered_inputs, gathered_slot_mapping
43 +
44 + if num_decode_tokens == local_num_tokens:
45 + return tensors, slot_mapping[:num_decode_tokens]
46 +
47 gathered_prefills = tuple(
48 pcp_group.all_gather(tensor[num_decode_tokens:].contiguous(), dim=0)
49 for tensor in tensors
50 )
36 - pcp_size = pcp_group.world_size
37 - gathered_slot_mapping = slot_mapping[: pcp_size * local_num_tokens]
51 if num_decode_tokens == 0:
52 return gathered_prefills, gathered_slot_mapping
53
@@ -58,6 +71,7 @@ def maybe_gather_mla_latent_cache_inputs(
71 slot_mapping: torch.Tensor | None,
72 num_decode_tokens: int | None,
73 use_pcp: bool,
74 + pcp_shard_decode_requests: bool = False,
75 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
76 if not use_pcp or num_decode_tokens is None:
77 return kv_c_normed, k_pe, slot_mapping
@@ -68,6 +82,7 @@ def maybe_gather_mla_latent_cache_inputs(
82 (kv_c_normed, k_pe_flat),
83 slot_mapping,
84 num_decode_tokens,
85 + pcp_shard_decode_requests,
86 )
87 cache_k_pe = cache_k_pe_flat.view(-1, *k_pe.shape[1:])
88 return cache_kv_c, cache_k_pe, cache_slot_mapping
@@ -78,11 +93,12 @@ def maybe_gather_indexer_k(
93 slot_mapping: torch.Tensor,
94 num_decode_tokens: int,
95 use_pcp: bool,
96 + pcp_shard_decode_requests: bool = False,
97 ) -> tuple[torch.Tensor, torch.Tensor]:
98 if not use_pcp:
99 return k, slot_mapping
100 (cache_k,), cache_slot_mapping = _gather_prefill_cache_inputs(
85 - (k,), slot_mapping, num_decode_tokens
101 + (k,), slot_mapping, num_decode_tokens, pcp_shard_decode_requests
102 )
103 return cache_k, cache_slot_mapping
104
vllm/v1/worker/gpu/pcp_manager.py
+26 -3
@@ -49,6 +49,7 @@ class PCPManager:
49 pcp_world_size: int,
50 pcp_rank: int,
51 device: torch.device,
52 + shard_decode_requests: bool,
53 max_num_reqs: int | None = None,
54 max_num_tokens: int | None = None,
55 block_tables: BlockTables | None = None,
@@ -62,6 +63,7 @@ class PCPManager:
63 self.dcp_world_size = dcp_world_size
64 self.dcp_rank = dcp_rank
65 self.cp_interleave = cp_interleave
66 + self.shard_decode_requests = shard_decode_requests
67
68 self._global_batch: InputBatch | None = None
69 self._local_batch: InputBatch | None = None
@@ -248,6 +250,7 @@ class PCPManager:
250
251 Decodes, and prefills too short to fill all eight, are replicated instead.
252 """
253 + decode_ordinal = 0
254 num_chunks = 2 * self.pcp_world_size
255 replicated = self.replicated_requests(num_scheduled_tokens, is_prefilling)
256 for global_batch_req_idx, num_tokens in enumerate(num_scheduled_tokens):
@@ -258,7 +261,16 @@ class PCPManager:
261 if not replicated[global_batch_req_idx]:
262 chunk_size = (query_len + num_chunks - 1) // num_chunks
263 chunk_indices = (rank, num_chunks - 1 - rank)
261 - else: # decodes, and short prefills under DCP, are replicated
264 + elif self.shard_decode_requests:
265 + chunk_size = query_len
266 + # KV and hidden states are gathered back to every PCP rank, so
267 + # decode ownership does not need to persist across steps. Use a
268 + # compact decode-only ordinal to keep each step exactly
269 + # balanced even when prefills and zero-token rows are present.
270 + owner_rank = decode_ordinal % self.pcp_world_size
271 + decode_ordinal += 1
272 + chunk_indices = (0,) if rank == owner_rank else ()
273 + else: # DCP requires decode queries on every participating rank.
274 chunk_size = query_len
275 chunk_indices = (0,)
276
@@ -346,8 +358,18 @@ class PCPManager:
358 segment.global_batch_slice.stop,
359 dtype=np.int64,
360 )
349 - # Cache insertion pairs one slot entry with each rank's local decode.
350 - if replicated[segment.global_batch_req_idx] and rank != 0:
361 + # Replicated rows contain identical KV, so rank 0 is the
362 + # canonical writer. A sharded decode has exactly one owner and
363 + # therefore needs no de-duplication here.
364 + is_sharded_decode = (
365 + not bool(is_prefilling[segment.global_batch_req_idx])
366 + and self.shard_decode_requests
367 + )
368 + if (
369 + replicated[segment.global_batch_req_idx]
370 + and not is_sharded_decode
371 + and rank != 0
372 + ):
373 continue
374 gathered_kv_write_mask[padded_gathered_slice] = True
375 hidden_restore_idx[segment.global_batch_slice] = np.arange(
@@ -829,6 +851,7 @@ def maybe_build_pcp_manager(
851 pcp_world_size=pcp_size,
852 pcp_rank=pcp_rank,
853 device=device,
854 + shard_decode_requests=parallel_config.pcp_shard_decode_requests,
855 max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
856 max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens,
857 block_tables=block_tables,