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

[PCP][DCP] Enable PCP+DCP on sparse-MLA models (#56157)

Signed-off-by: patryk.saffer <patryk.saffer@blc-login-1.blc1.mistralcompute.ai> Signed-off-by: patryk.saffer <patryk.saffer@blc-login-0.blc1.mistralcompute.ai> Signed-off-by: PatrykSaffer <patryk.saffer@mistral.ai> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: patryk.saffer <patryk.saffer@blc-login-1.blc1.mistralcompute.ai> Co-authored-by: patryk.saffer <patryk.saffer@blc-login-0.blc1.mistralcompute.ai> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Patryk Saffer <PatrykSaffer@users.noreply.github.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

PatrykSaffer committed Sep 13, 2026 at 01:24 UTC ebe1dec2da16ac559e40446a4c6e06a73c9efe00
21 files changed +917 -138
.buildkite/test_areas/lm_eval.yaml
+5
@@ -110,12 +110,17 @@ steps:
110 - csrc/
111 - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
112 - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
113 + - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml
114 - tests/evals/gsm8k/configs/models-pcp.txt
115 - vllm/model_executor/layers/quantization
116 - vllm/config/parallel.py
117 - vllm/distributed/parallel_state.py
118 - vllm/model_executor/layers/attention/mla_attention.py
119 - vllm/model_executor/layers/attention/pcp.py
120 + - vllm/model_executor/layers/attention/sparse_mla_attention.py
121 + - vllm/model_executor/layers/sparse_attn_indexer.py
122 + - vllm/v1/attention/backends/mla/flashmla_sparse.py
123 + - vllm/v1/attention/backends/mla/indexer.py
124 - vllm/v1/worker/gpu/model_runner.py
125 - vllm/v1/worker/gpu/pcp_manager.py
126 autorun_on_main: true
tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml new
+19
@@ -0,0 +1,19 @@
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 + --compilation-config '{"cudagraph_mode":"PIECEWISE"}'
8 + --max-model-len 4096
9 + --max-num-batched-tokens 32768
10 + --safetensors-load-strategy prefetch
11 + --moe-backend flashinfer_cutlass
12 + --prefill-context-parallel-size 4
13 + --decode-context-parallel-size 4
14 + --enable-expert-parallel
15 + --kv-cache-dtype fp8
16 +env:
17 + PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True"
18 + VLLM_LOGGING_LEVEL: "DEBUG"
19 + VLLM_USE_V2_MODEL_RUNNER: "1"
tests/evals/gsm8k/configs/models-pcp.txt
+1
@@ -1,2 +1,3 @@
1 GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
2 GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
3 +GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml
tests/v1/attention/test_indexer_dcp_localize.py
+35
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import numpy as np
5 import pytest
6 import torch
7
@@ -962,3 +963,37 @@ def test_sparse_decode_dcp_short_context_matches_non_dcp():
963 dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses)
964 torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5)
965 torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5)
966 +
967 +
968 +@pytest.mark.parametrize("dcp_world_size", [2, 4, 8])
969 +@pytest.mark.parametrize("req_lens", [[7], [8, 8], [1, 9], [256, 1, 2730]])
970 +def test_pcp_plan_deinterleave_restores_global_order(dcp_world_size, req_lens):
971 + """The index gather must undo DCP sharding per request"""
972 + from vllm.v1.attention.backends.mla.indexer import build_pcp_global_chunk_plan
973 +
974 + scheduled = np.array(req_lens, dtype=np.int64)
975 + shard_rows = -(-scheduled // dcp_world_size)
976 + rows = np.arange(len(scheduled))
977 + plan = build_pcp_global_chunk_plan(
978 + rows, shard_rows, dcp_world_size, torch.device("cpu")
979 + )
980 +
981 + starts = np.concatenate([[0], np.cumsum(scheduled)])
982 +
983 + # Build each rank's padded shard exactly as the cache gather would.
984 + padded_cu = plan.padded_local_cu.tolist()
985 + shards = torch.zeros(dcp_world_size, plan.padded_local_total, 1)
986 + for r in range(dcp_world_size):
987 + for i, g in enumerate(req_lens):
988 + for t in range(r, g, dcp_world_size):
989 + shards[r, padded_cu[i] + t // dcp_world_size, 0] = starts[i] + t
990 +
991 + gathered = shards.reshape(dcp_world_size * plan.padded_local_total, 1)
992 + restored = gathered[plan.deinterleave_idx]
993 + # The layout is padded per request; only the real positions are read.
994 + row_start = plan.row_start_cu.tolist()
995 + for i, g in enumerate(req_lens):
996 + torch.testing.assert_close(
997 + restored[row_start[i] : row_start[i] + g, 0],
998 + torch.arange(starts[i], starts[i] + g, dtype=torch.float32),
999 + )
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+1
@@ -212,6 +212,7 @@ def test_index_conversion_warmup_uses_physical_block_stride():
212 ),
213 parallel_config=SimpleNamespace(
214 decode_context_parallel_size=1,
215 + prefill_context_parallel_size=1,
216 cp_kv_cache_interleave_size=1,
217 ),
218 )
tests/v1/attention/test_sparse_mla_backends.py
+3
@@ -1080,6 +1080,7 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens():
1080 _convert_logical_to_physical_topk=_convert_topk,
1081 index_group=None,
1082 index_group_index=0,
1083 + pcp_dcp_kv_gather=False,
1084 )
1085
1086 out, _ = FlashMLASparseImpl._forward_bf16_kv(
@@ -3257,6 +3258,7 @@ def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch):
3258 device=torch.device(DEVICE_TYPE),
3259 dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE),
3260 max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE),
3261 + pcp_dcp_kv_gather=False,
3262 )
3263 query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660])
3264 common_metadata = SimpleNamespace(
@@ -3363,6 +3365,7 @@ def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: b
3365 index_group=None,
3366 index_group_index=0,
3367 dcp_world_size=1,
3368 + pcp_dcp_kv_gather=False,
3369 need_to_return_lse_for_decode=False,
3370 _fp8_flash_mla_kernel=run_kernel,
3371 _convert_logical_to_physical_topk=(
tests/v1/worker/test_gpu_pcp_manager.py
+149
@@ -114,6 +114,155 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
114 )
115
116
117 +def _rank_rows(
118 + pcp_rank: int,
119 + pcp_world_size: int,
120 + num_scheduled_tokens: np.ndarray,
121 + num_computed_tokens: np.ndarray,
122 + is_prefilling: np.ndarray,
123 +) -> tuple[np.ndarray, np.ndarray]:
124 + """One PCP+DCP rank's rows as (global request, extent)."""
125 + manager = PCPManager(
126 + pcp_world_size=pcp_world_size,
127 + pcp_rank=pcp_rank,
128 + device=torch.device("cpu"),
129 + dcp_world_size=2,
130 + )
131 + query_start_loc_np = np.concatenate([[0], np.cumsum(num_scheduled_tokens)]).astype(
132 + np.int32
133 + )
134 + segments = manager._get_rank_segments(
135 + pcp_rank, num_scheduled_tokens, is_prefilling, query_start_loc_np
136 + )
137 + rows = np.array([segment.global_batch_req_idx for segment in segments])
138 + extents = (num_computed_tokens + num_scheduled_tokens)[rows]
139 + return rows, extents
140 +
141 +
142 +@pytest.mark.parametrize(
143 + ("num_scheduled_tokens", "num_computed_tokens"),
144 + [
145 + ([32, 32], [0, 0]),
146 + ([32, 32, 32, 32], [0, 0, 0, 0]),
147 + ([32, 32], [100, 0]), # one continued, one fresh
148 + ([64, 32, 48], [0, 512, 0]), # ragged lengths and mixed contexts
149 + # A replicated prefill BEHIND a split one. Sorting replicated rows first
150 + # put request 1 ahead of request 0.
151 + ([64, 3], [0, 0]),
152 + ([64, 3, 64], [0, 0, 0]),
153 + ],
154 +)
155 +@pytest.mark.parametrize("pcp_world_size", [2, 4])
156 +def test_published_row_order_is_identical_on_every_pcp_rank(
157 + num_scheduled_tokens, num_computed_tokens, pcp_world_size
158 +):
159 + """Every rank must map its rows to the same global requests, in the same order."""
160 + num_scheduled_tokens = np.array(num_scheduled_tokens, dtype=np.int32)
161 + num_computed_tokens = np.array(num_computed_tokens, dtype=np.int32)
162 + is_prefilling = np.ones(len(num_scheduled_tokens), dtype=np.bool_)
163 +
164 + orders = [
165 + _rank_rows(
166 + rank,
167 + pcp_world_size,
168 + num_scheduled_tokens,
169 + num_computed_tokens,
170 + is_prefilling,
171 + )[0]
172 + for rank in range(pcp_world_size)
173 + ]
174 + for rank, order in enumerate(orders[1:], start=1):
175 + assert np.array_equal(orders[0], order), (
176 + f"rank {rank} rows map to {order.tolist()}, rank 0 to {orders[0].tolist()}"
177 + )
178 + # Grouped by request and ascending: what the indexer plan indexes by.
179 + assert np.all(np.diff(orders[0]) >= 0)
180 +
181 +
182 +def test_published_row_order_puts_every_decode_before_every_prefill():
183 + """split_decodes_and_prefills takes the FIRST prefilling row as the boundary."""
184 + # req 0 is a continued prefill, req 1 is a decode.
185 + req_idx, _ = _rank_rows(
186 + pcp_rank=0,
187 + pcp_world_size=2,
188 + num_scheduled_tokens=np.array([32, 1], dtype=np.int32),
189 + num_computed_tokens=np.array([100, 20], dtype=np.int32),
190 + is_prefilling=np.array([True, False], dtype=np.bool_),
191 + )
192 + # The decode (request 1) must be row 0.
193 + assert req_idx[0] == 1
194 +
195 +
196 +def test_split_prefill_rows_repeat_the_request_and_its_full_extent():
197 + """Equal adjacent request indices are what let a backend share a KV region."""
198 + req_idx, extents = _rank_rows(
199 + pcp_rank=1,
200 + pcp_world_size=2,
201 + num_scheduled_tokens=np.array([64, 3], dtype=np.int32),
202 + num_computed_tokens=np.array([100, 0], dtype=np.int32),
203 + is_prefilling=np.ones(2, dtype=np.bool_),
204 + )
205 + assert req_idx.tolist() == [0, 0, 1]
206 + assert extents.tolist() == [164, 164, 3]
207 +
208 +
209 +@pytest.mark.parametrize("query_len", [1, 2, 3, 5, 6, 9])
210 +def test_dcp_replicates_prefills_too_short_to_split(query_len):
211 + """A prefill that cannot fill 2*pcp chunks is replicated, not split."""
212 + pcp_world_size = 2
213 + num_scheduled_tokens = np.array([query_len], dtype=np.int32)
214 + is_prefilling = np.ones(1, dtype=np.bool_)
215 +
216 + rows_per_rank = []
217 + for rank in range(pcp_world_size):
218 + manager = PCPManager(
219 + pcp_world_size=pcp_world_size,
220 + pcp_rank=rank,
221 + device=torch.device("cpu"),
222 + dcp_world_size=2,
223 + )
224 + assert manager.replicated_requests(num_scheduled_tokens, is_prefilling)[0]
225 + rows = list(
226 + manager._iter_rank_chunks(rank, num_scheduled_tokens, is_prefilling)
227 + )
228 + assert rows == [(0, 0, query_len)]
229 + rows_per_rank.append(rows)
230 + assert rows_per_rank[0] == rows_per_rank[1]
231 +
232 +
233 +@pytest.mark.parametrize("pcp_world_size", [2, 4, 8])
234 +@pytest.mark.parametrize(
235 + "query_len", [16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1000, 4097]
236 +)
237 +def test_pcp_first_chunk_row_is_never_short(pcp_world_size, query_len):
238 + num_scheduled_tokens = np.array([query_len], dtype=np.int32)
239 + is_prefilling = np.ones(1, dtype=np.bool_)
240 +
241 + charged_per_rank = []
242 + for rank in range(pcp_world_size):
243 + manager = PCPManager(
244 + pcp_world_size=pcp_world_size,
245 + pcp_rank=rank,
246 + device=torch.device("cpu"),
247 + dcp_world_size=pcp_world_size,
248 + )
249 + chunk_lens = [
250 + chunk_len
251 + for _, _, chunk_len in manager._iter_rank_chunks(
252 + rank, num_scheduled_tokens, is_prefilling
253 + )
254 + ]
255 + assert chunk_lens, f"rank {rank} got no rows for {query_len=}"
256 + assert chunk_lens[0] == max(chunk_lens), (
257 + f"rank {rank} emitted a short first chunk for {query_len=}: {chunk_lens}"
258 + )
259 + charged_per_rank.append(len(chunk_lens) * chunk_lens[0])
260 +
261 + assert len(set(charged_per_rank)) == 1, (
262 + f"ranks would chunk differently for {query_len=}: {charged_per_rank}"
263 + )
264 +
265 +
266 def _make_global_decode_batch(
267 num_computed_tokens: list[int], buffers: InputBuffers, device: torch.device
268 ) -> InputBatch:
vllm/model_executor/layers/attention/mla_attention.py
+7
@@ -420,6 +420,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):
420 """
421
422 supports_dense_mha_prefill: ClassVar[bool] = True
423 + # Under PCP+DCP only the decode rows carry an LSE; the base forward
424 + # merges a full-batch LSE, so subclasses opt in with their own forward.
425 + supports_pcp_dcp: ClassVar[bool] = False
426
427 def __init__(
428 self,
@@ -660,6 +663,10 @@ class MLAAttention(nn.Module, AttentionLayerBase):
663
664 _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()
665
666 + if self.use_pcp and self.impl.dcp_world_size > 1 and not self.supports_pcp_dcp:
667 + raise NotImplementedError(
668 + f"{type(self).__name__} does not support PCP+DCP."
669 + )
670 self.dcp_manager: MLADCPManager | None = None
671 if self.impl.dcp_world_size > 1:
672 query_dtype = (
vllm/model_executor/layers/attention/sparse_mla_attention.py
+8 -3
@@ -190,7 +190,8 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
190 pin_memory=PIN_MEMORY,
191 )
192 parallel_config = vllm_config.parallel_config
193 - self.use_pcp = parallel_config.prefill_context_parallel_size > 1
193 + self.pcp_world_size = parallel_config.prefill_context_parallel_size
194 + self.use_pcp = self.pcp_world_size > 1
195 try:
196 self.dcp_world_size = get_dcp_group().world_size
197 except AssertionError:
@@ -409,8 +410,12 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
410 q_data_type=self.model_config.dtype,
411 output_dtype=self.model_config.dtype,
412 prefill_backend=self._prefill_backend,
412 - use_dense_mha=_use_dense_mha_prefill(
413 - self.vllm_config, prefill_max_seq_len, self.topk_tokens
413 + use_dense_mha=(
414 + _use_dense_mha_prefill(
415 + self.vllm_config, prefill_max_seq_len, self.topk_tokens
416 + )
417 + # Dense MHA cannot be used under PCP+DCP.
418 + and not (self.use_pcp and self.dcp_world_size > 1)
419 ),
420 topk_mask_workspace=self.topk_mask_workspace,
421 host_staging_plan=staging_plan,
vllm/model_executor/layers/sparse_attn_indexer.py
+81 -13
@@ -3,6 +3,7 @@
3 """Custom Sparse Attention Indexer layers."""
4
5 import torch
6 +import torch.distributed as dist
7
8 import vllm.envs as envs
9 from vllm import _custom_ops as ops
@@ -130,6 +131,23 @@ def _merge_dcp_topk_global(
131 )
132
133
134 +def dcp_gather_kv_rows(
135 + local_padded: torch.Tensor,
136 + deinterleave_idx: torch.Tensor,
137 + gathered_buf: torch.Tensor,
138 + out_buf: torch.Tensor,
139 +) -> torch.Tensor:
140 + """All-gather this rank's KV shard and return it in global token order."""
141 +
142 + gathered = gathered_buf[: get_dcp_group().world_size * local_padded.shape[0]]
143 + dist.all_gather_into_tensor(
144 + gathered, local_padded.contiguous(), group=get_dcp_group().device_group
145 + )
146 + return torch.index_select(
147 + gathered, 0, deinterleave_idx, out=out_buf[: deinterleave_idx.shape[0]]
148 + )
149 +
150 +
151 @triton.jit
152 def _fused_indexer_q_rope_quant_kernel(
153 positions,
@@ -346,11 +364,19 @@ def sparse_attn_indexer(
364 values_spec, scales_spec = _gather_workspace_shapes(
365 total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
366 )
349 - current_workspace_manager().get_simultaneous(
367 + profile_specs: list[tuple[tuple[int, ...], torch.dtype]] = [
368 values_spec,
369 scales_spec,
370 ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
353 - )
371 + ]
372 + if use_pcp and dcp_world_size > 1:
373 + # The PCP+DCP path takes an all-gather destination and a
374 + # de-interleaved result.
375 + gather_spec = _gather_workspace_shapes(
376 + total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
377 + )
378 + profile_specs.extend(gather_spec * 2)
379 + current_workspace_manager().get_simultaneous(*profile_specs)
380
381 # Dummy allocation to simulate for peak logits tensor memory during inference.
382 # FP8 elements so elements == bytes
@@ -462,9 +488,27 @@ def sparse_attn_indexer(
488 values_spec, scales_spec = _gather_workspace_shapes(
489 total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
490 )
465 - k_quant_full, k_scale_full = workspace_manager.get_simultaneous(
491 + # PCP + DCP needs two more pairs: the rank-major all-gather destination
492 + # and the de-interleaved result.
493 + pcp_chunks = [
494 + (c, c.pcp_deinterleave_idx)
495 + for c in prefill_metadata.chunks
496 + if c.pcp_deinterleave_idx is not None
497 + ]
498 + gather_specs: list[tuple[tuple[int, int], torch.dtype]] = []
499 + if pcp_chunks:
500 + gathered_rows = max(
501 + dcp_world_size * c.local_total_seq_lens for c, _ in pcp_chunks
502 + )
503 + deinterleaved_rows = max(idx.shape[0] for _, idx in pcp_chunks)
504 + for rows in (gathered_rows, deinterleaved_rows):
505 + gather_specs.extend(
506 + _gather_workspace_shapes(rows, head_dim, fp8_dtype, use_fp4_cache)
507 + )
508 + k_quant_full, k_scale_full, *gather_bufs = workspace_manager.get_simultaneous(
509 values_spec,
510 scales_spec,
511 + *gather_specs,
512 )
513 for chunk in prefill_metadata.chunks:
514 cu_seqlen_ks = chunk.cu_seqlen_ks
@@ -481,6 +525,23 @@ def sparse_attn_indexer(
525 chunk.local_cu_seq_lens,
526 )
527
528 + # PCP + DCP KV all-gather.
529 + deinterleave_idx = chunk.pcp_deinterleave_idx
530 + if deinterleave_idx is not None:
531 + # local_total_seq_lens is the PCP-padded extent here, so every
532 + # rank contributes an identically shaped shard.
533 + gathered_values, gathered_scales, out_values, out_scales = gather_bufs
534 + shard = k_quant[: chunk.local_total_seq_lens]
535 + k_quant = dcp_gather_kv_rows(
536 + shard, deinterleave_idx, gathered_values, out_values
537 + )
538 + k_scale = dcp_gather_kv_rows(
539 + k_scale[: chunk.local_total_seq_lens],
540 + deinterleave_idx,
541 + gathered_scales,
542 + out_scales,
543 + )
544 +
545 q_slice = q_quant[chunk.token_start : chunk.token_end]
546 q_scale_slice = (
547 q_scale[chunk.token_start : chunk.token_end]
@@ -491,7 +552,7 @@ def sparse_attn_indexer(
552 chunk.token_start : chunk.token_end, :topk_tokens
553 ]
554
494 - if chunk.local_total_seq_lens == 0:
555 + if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0:
556 logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32)
557 topk_indices.fill_(-1)
558 else:
@@ -561,15 +622,18 @@ def sparse_attn_indexer(
622 topk_tokens,
623 )
624
564 - _merge_dcp_topk_global(
565 - logits,
566 - topk_indices,
567 - topk_tokens,
568 - dcp_rank,
569 - dcp_world_size,
570 - cp_kv_cache_interleave_size,
571 - row_starts=chunk.cu_seqlen_ks,
572 - )
625 + if deinterleave_idx is None:
626 + # Under the PCP path the top-k already ran over the whole
627 + # context.
628 + _merge_dcp_topk_global(
629 + logits,
630 + topk_indices,
631 + topk_tokens,
632 + dcp_rank,
633 + dcp_world_size,
634 + cp_kv_cache_interleave_size,
635 + row_starts=chunk.cu_seqlen_ks,
636 + )
637
638 if has_decode:
639 decode_metadata = attn_metadata_narrowed.decode
@@ -910,6 +974,10 @@ class SparseAttnIndexer(CustomOp):
974 if current_platform.is_cuda() or current_platform.is_xpu():
975 return self.forward_cuda(hidden_states, q_quant, k, weights)
976 elif current_platform.is_rocm():
977 + if self.use_pcp and self.dcp_world_size > 1:
978 + raise NotImplementedError(
979 + "The ROCm sparse-indexer path does not support PCP+DCP."
980 + )
981 return self.forward_hip(hidden_states, q_quant, k, weights)
982 elif current_platform.is_cpu():
983 return self.forward_cpu(hidden_states, q_quant, k, weights)
vllm/model_executor/layers/sparse_attn_indexer_kpool.py
+10
@@ -919,6 +919,16 @@ class SparseAttnIndexerKpool(CustomOp):
919 raise RuntimeError(
920 "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed."
921 )
922 + _cfg = get_current_vllm_config_or_none()
923 + _parallel = _cfg.parallel_config if _cfg is not None else None
924 + if (
925 + _parallel is not None
926 + and _parallel.prefill_context_parallel_size > 1
927 + and _parallel.decode_context_parallel_size > 1
928 + ):
929 + raise NotImplementedError(
930 + "SparseAttnIndexerKpool does not support PCP+DCP."
931 + )
932
933 def forward_native(
934 self,
vllm/models/deepseek_v32/attention.py
+22 -9
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import TYPE_CHECKING, cast
3 +from typing import TYPE_CHECKING, ClassVar, cast
4
5 import torch
6 import torch.nn as nn
@@ -121,6 +121,7 @@ class DeepseekV32Indexer(nn.Module):
121 class DeepseekV32Attention(MLAAttention):
122 indexer: "DeepseekV32Indexer | None"
123 indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer
124 + supports_pcp_dcp: ClassVar[bool] = True
125
126 def __init__(
127 self,
@@ -569,8 +570,9 @@ class DeepseekV32Attention(MLAAttention):
570 seq_lens: torch.Tensor | None
571 query_start_loc: torch.Tensor | None
572 if self.use_pcp:
572 - if attn_metadata.decode is not None:
573 - seq_lens = attn_metadata.decode.seq_lens
573 + decode_metadata = getattr(attn_metadata, "decode", None)
574 + if decode_metadata is not None:
575 + seq_lens = decode_metadata.seq_lens
576 else:
577 all_seq_lens = cast(
578 torch.Tensor,
@@ -585,12 +587,23 @@ class DeepseekV32Attention(MLAAttention):
587 # PCP-only empty-shard metadata is needed.
588 seq_lens = None
589 query_start_loc = None
588 - attn_out = self.dcp_manager.combine(
589 - attn_out,
590 - lse,
591 - seq_lens=seq_lens, # type: ignore[arg-type]
592 - query_start_loc=query_start_loc, # type: ignore[arg-type]
593 - )
590 + # Under PCP+DCP the prefill rows attended over the gathered KV, so
591 + # only the decode rows carry an LSE and take part in the merge.
592 + num_merge_rows = lse.shape[0]
593 + if num_merge_rows == attn_out.shape[0]:
594 + attn_out = self.dcp_manager.combine(
595 + attn_out,
596 + lse,
597 + seq_lens=seq_lens, # type: ignore[arg-type]
598 + query_start_loc=query_start_loc, # type: ignore[arg-type]
599 + )
600 + elif num_merge_rows > 0:
601 + attn_out[:num_merge_rows] = self.dcp_manager.combine(
602 + attn_out[:num_merge_rows],
603 + lse,
604 + seq_lens=seq_lens, # type: ignore[arg-type]
605 + query_start_loc=query_start_loc, # type: ignore[arg-type]
606 + )
607 if self.use_pcp:
608 attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads)
609
vllm/v1/attention/backend.py
+15 -2
@@ -415,9 +415,14 @@ class CommonAttentionMetadata:
415 encoder_seq_lens_cpu: np.ndarray | None = None
416
417 dcp_local_seq_lens: torch.Tensor | None = None
418 - dcp_local_seq_lens_cpu: torch.Tensor | None = None
418 """Sequence lengths of the local rank in decode context parallelism world"""
419
420 + dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None
421 + """(batch_size,) CPU upper bound on dcp_local_seq_lens. Under PCP+DCP it
422 + holds, on every row, the largest shard any DCP rank has of the row's
423 + whole request, identical on every PCP rank, so the sparse backends pad
424 + their KV gather to it."""
425 +
426 positions: torch.Tensor | None = None
427 """(num_actual_tokens,) token positions. Optional; set when the caller
428 has positions available so that builders can pre-compute position-dependent
@@ -434,6 +439,11 @@ class CommonAttentionMetadata:
439 decode rows (assumes every draft was accepted). Not safe for kernels
440 that need exact per-row context lengths on decode rows."""
441
442 + req_idx: np.ndarray | None = None
443 + """(batch_size,) index of each row's request in the runner's request
444 + table. Rows of one request are adjacent, so equal neighbours are PCP
445 + chunks sharing one KV context."""
446 +
447 mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None
448 """PrefixLM bidirectional ranges for multimodal tokens. Maps
449 request index to list of (start, end) token position ranges
@@ -525,9 +535,12 @@ class CommonAttentionMetadata:
535 encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens),
536 encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu),
537 dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),
528 - dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu),
538 + dcp_local_seq_lens_cpu_upper_bound=maybe_slice_reqs(
539 + self.dcp_local_seq_lens_cpu_upper_bound
540 + ),
541 seq_lens_cpu_upper_bound=maybe_slice_reqs(self.seq_lens_cpu_upper_bound),
542 is_prefilling=maybe_slice_reqs(self.is_prefilling),
543 + req_idx=maybe_slice_reqs(self.req_idx),
544 rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens),
545 replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu),
546 )
vllm/v1/attention/backends/mla/flashmla_sparse.py
+215 -57
@@ -3,11 +3,14 @@
3 from dataclasses import dataclass
4 from typing import TYPE_CHECKING, ClassVar
5
6 +import numpy as np
7 import torch
8 +import torch.distributed as dist
9
10 from vllm import _custom_ops as ops
11 from vllm.config import VllmConfig, get_current_vllm_config
12 from vllm.config.cache import CacheDType
13 +from vllm.distributed.parallel_state import get_dcp_group
14 from vllm.logger import init_logger
15 from vllm.model_executor.layers.attention.sparse_mla_attention import (
16 SparseMLACommonImpl,
@@ -28,10 +31,12 @@ from vllm.v1.attention.backend import (
31 from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
32 from vllm.v1.attention.backends.mla.sparse_utils import (
33 flat_kv_row_view,
34 + request_row_bounds,
35 triton_convert_req_index_to_global_index,
36 triton_filter_and_convert_dcp_index,
37 )
38 from vllm.v1.attention.backends.utils import (
39 + get_dcp_local_seq_lens,
40 reshape_attn_output_for_spec_decode,
41 reshape_query_for_spec_decode,
42 split_prefill_chunks,
@@ -43,6 +48,7 @@ from vllm.v1.attention.ops.flashmla import (
48 get_mla_metadata,
49 )
50 from vllm.v1.kv_cache_interface import AttentionSpec
51 +from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
52 from vllm.v1.worker.workspace import current_workspace_manager
53
54 if TYPE_CHECKING:
@@ -171,6 +177,28 @@ class FlashMLASparseBackend(AttentionBackend):
177 return None
178
179
180 +def gathered_prefill_shards(
181 + row_req_idx: np.ndarray, row_shard_rows: np.ndarray
182 +) -> tuple[np.ndarray, np.ndarray]:
183 + """Under PCP+DCP, group this rank's prefill rows by request and take each
184 + request's slice of this rank's KV shard.
185 +
186 + PCP gives a rank two adjacent chunk rows of a split prefill; they share
187 + one context, so the workspace holds it once. ``row_shard_rows`` is
188 + ``CommonAttentionMetadata.dcp_local_seq_lens_cpu_upper_bound``: the
189 + largest DCP shard of the whole request on every row, identical on every
190 + PCP rank, so the all-gather is shaped the same everywhere.
191 +
192 + Returns the request row bounds and, per request, its shard rows.
193 + """
194 + row_bounds = request_row_bounds(row_req_idx)
195 + shard_rows = row_shard_rows[row_bounds[:-1]].astype(np.int64)
196 + assert np.all(shard_rows > 0), (
197 + f"PCP+DCP prefill got an empty context: {shard_rows.tolist()}"
198 + )
199 + return row_bounds, shard_rows
200 +
201 +
202 @dataclass
203 class FlashMLASparseMetadata(SparseMLACommonMetadata):
204 @dataclass
@@ -331,8 +359,13 @@ class FlashMLASparseMetadataBuilder(
359 device=device,
360 )
361
362 + # PCP+DCP attends prefill rows over the DCP-gathered KV, which only the
363 + # separate prefill/decode path can do.
364 + self.pcp_dcp_kv_gather = self.use_pcp and self.dcp_world_size > 1
365 self.fp8_use_mixed_batch = (
335 - self.num_heads < MIN_HEADS_FOR_BF16_PREFILL and not self.use_hisparse
366 + self.num_heads < MIN_HEADS_FOR_BF16_PREFILL
367 + and not self.use_hisparse
368 + and not self.pcp_dcp_kv_gather
369 )
370
371 if parallel_config.decode_context_parallel_size > 1:
@@ -342,7 +375,12 @@ class FlashMLASparseMetadataBuilder(
375 "default 'ag_rs' DCP comm backend; got "
376 f"'{parallel_config.dcp_comm_backend}'"
377 )
345 - if not self.fp8_use_mixed_batch:
378 + if self.pcp_dcp_kv_gather and cache_config.cache_dtype != "fp8_ds_mla":
379 + raise NotImplementedError(
380 + "PCP+DCP sparse prefill gathers the KV through the fp8_ds_mla "
381 + f"upconvert; got a {cache_config.cache_dtype} cache"
382 + )
383 + if not self.fp8_use_mixed_batch and not self.pcp_dcp_kv_gather:
384 raise NotImplementedError(
385 "DCP for FlashMLA sparse is only supported on the "
386 "mixed-batch fp8 path (num_heads < "
@@ -353,9 +391,16 @@ class FlashMLASparseMetadataBuilder(
391 # Head padding (and the tile-scheduler metadata sized from it) is
392 # computed from the local head count, but the kernel runs on the
393 # DCP-gathered heads.
356 - gathered_num_heads = (
357 - self.num_heads * parallel_config.decode_context_parallel_size
358 - )
394 + if self.use_pcp:
395 + gathered_num_heads = (
396 + self.num_heads * parallel_config.tensor_parallel_size
397 + if self.dcp_world_size > self.pcp_world_size
398 + else self.num_heads
399 + )
400 + else:
401 + gathered_num_heads = (
402 + self.num_heads * parallel_config.decode_context_parallel_size
403 + )
404 gathered_padded_heads = FlashMLASparseImpl._compute_fp8_decode_padded_heads(
405 gathered_num_heads
406 )
@@ -433,6 +478,40 @@ class FlashMLASparseMetadataBuilder(
478 query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
479
480 prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:]
481 + # Chunk prefill requests to fit within workspace size
482 + max_prefill_buffer_size = get_prefill_workspace_size(
483 + self.vllm_config.model_config.max_model_len
484 + )
485 + # Workspace rows per entry, and the prefill rows each entry covers.
486 + workspace_rows = prefill_seq_lens_cpu
487 + row_bounds = np.arange(num_prefills + 1)
488 + prefill_block_table = common_attn_metadata.block_table_tensor[num_decodes:]
489 + prefill_seq_lens = common_attn_metadata.seq_lens[num_decodes:]
490 + if self.pcp_dcp_kv_gather:
491 + # One entry per request, holding this rank's KV shard: the
492 + # all-gather of every rank's shard is the whole context.
493 + # A dummy batch bypasses the PCP manager: one row per request,
494 + # so the row's own extent is the request's.
495 + row_req_idx = common_attn_metadata.req_idx
496 + if row_req_idx is None:
497 + row_req_idx = np.arange(common_attn_metadata.num_reqs)
498 + shard_rows_cpu = common_attn_metadata.dcp_local_seq_lens_cpu_upper_bound
499 + if shard_rows_cpu is None:
500 + shard_rows_cpu = get_dcp_local_seq_lens(
501 + seq_lens_cpu,
502 + self.dcp_world_size,
503 + 0,
504 + self.cp_kv_cache_interleave_size,
505 + )
506 + row_bounds, rows_per_rank = gathered_prefill_shards(
507 + row_req_idx[num_decodes:], shard_rows_cpu[num_decodes:].numpy()
508 + )
509 + workspace_rows = torch.from_numpy(rows_per_rank.astype(np.int32))
510 + max_prefill_buffer_size //= self.dcp_world_size
511 + entry_rows = async_copy_to_gpu(row_bounds[:-1], device=self.device)
512 + prefill_block_table = prefill_block_table.index_select(0, entry_rows)
513 + prefill_seq_lens = prefill_seq_lens.index_select(0, entry_rows)
514 + num_entries = len(workspace_rows)
515
516 # Build prefill_request_id: -1 for decode, request index for
517 # prefill. This enables a single
@@ -441,33 +520,28 @@ class FlashMLASparseMetadataBuilder(
520 (num_tokens,), -1, dtype=torch.int32, device=self.device
521 )
522 # Map prefill tokens to their request IDs (0, 1, 2, ...)
444 - for req_idx in range(num_prefills):
523 + for req_idx in range(num_entries):
524 # Get query token range for this prefill request
446 - global_req_idx = num_decodes + req_idx
447 - req_query_start = query_start_loc_cpu[global_req_idx]
448 - req_query_end = query_start_loc_cpu[global_req_idx + 1]
525 + req_query_start = query_start_loc_cpu[
526 + num_decodes + int(row_bounds[req_idx])
527 + ]
528 + req_query_end = query_start_loc_cpu[
529 + num_decodes + int(row_bounds[req_idx + 1])
530 + ]
531 prefill_request_id[req_query_start:req_query_end] = req_idx
532
533 # will be adjusted by chunk loop
534 prefill_workspace_starts_cpu = torch.zeros(
453 - num_prefills, dtype=torch.int32, pin_memory=True
454 - )
455 - prefill_workspace_starts_cpu[1:] = torch.cumsum(
456 - prefill_seq_lens_cpu[:-1], dim=0
535 + num_entries, dtype=torch.int32, pin_memory=True
536 )
537 + prefill_workspace_starts_cpu[1:] = torch.cumsum(workspace_rows[:-1], dim=0)
538 # populated by non-blocking copy after prefill_workspace_starts_cpu is
539 # updated by each chunk
540 prefill_workspace_starts = torch.empty(
461 - num_prefills, dtype=torch.int32, device=self.device
541 + num_entries, dtype=torch.int32, device=self.device
542 )
543
464 - # Chunk prefill requests to fit within workspace size
465 - max_prefill_buffer_size = get_prefill_workspace_size(
466 - self.vllm_config.model_config.max_model_len
467 - )
468 - chunk_bounds = split_prefill_chunks(
469 - prefill_seq_lens_cpu, max_prefill_buffer_size
470 - )
544 + chunk_bounds = split_prefill_chunks(workspace_rows, max_prefill_buffer_size)
545
546 prefill_chunks = []
547 for chunk_start, chunk_end in chunk_bounds:
@@ -480,19 +554,19 @@ class FlashMLASparseMetadataBuilder(
554 offset = prefill_workspace_starts_cpu[chunk_start].item()
555 prefill_workspace_starts_cpu[chunk_start:chunk_end] -= offset
556
483 - chunk_tot_seqlen = prefill_seq_lens_cpu[chunk_start:chunk_end].sum()
484 - token_start = query_start_loc_cpu[num_decodes + chunk_start].item()
485 - token_end = query_start_loc_cpu[num_decodes + chunk_end].item()
557 + chunk_tot_seqlen = workspace_rows[chunk_start:chunk_end].sum()
558 + token_start = query_start_loc_cpu[
559 + num_decodes + int(row_bounds[chunk_start])
560 + ].item()
561 + token_end = query_start_loc_cpu[
562 + num_decodes + int(row_bounds[chunk_end])
563 + ].item()
564 tokens_slice = slice(token_start, token_end)
565
566 # Create chunk view of gpu tensor
567 chunk_workspace_starts = prefill_workspace_starts[chunk_start:chunk_end]
490 - chunk_block_table = common_attn_metadata.block_table_tensor[
491 - num_decodes + chunk_start : num_decodes + chunk_end
492 - ]
493 - chunk_seq_lens = common_attn_metadata.seq_lens[
494 - num_decodes + chunk_start : num_decodes + chunk_end
495 - ]
568 + chunk_block_table = prefill_block_table[chunk_start:chunk_end]
569 + chunk_seq_lens = prefill_seq_lens[chunk_start:chunk_end]
570
571 prefill_chunks.append(
572 FP8Meta.Prefill.Chunk(
@@ -516,14 +590,19 @@ class FlashMLASparseMetadataBuilder(
590 )
591
592 if num_decodes > 0:
519 - # Use padded head count since that's what the kernel will see
520 - scheduler_metadata, _ = get_mla_metadata()
593 + if self.pcp_dcp_kv_gather:
594 + # Decode rows take the mixed-batch kernel call, which returns
595 + # the LSE the DCP merge needs.
596 + kernel_meta = self._build_fp8_mixed_decode_prefill()
597 + else:
598 + # Use padded head count since that's what the kernel will see
599 + scheduler_metadata, _ = get_mla_metadata()
600
522 - kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata(
523 - scheduler_metadata=scheduler_metadata,
524 - dummy_block_table=self.dummy_block_table[:active_num_decodes],
525 - cache_lens=self.max_model_len_tensor[:active_num_decodes],
526 - )
601 + kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata(
602 + scheduler_metadata=scheduler_metadata,
603 + dummy_block_table=self.dummy_block_table[:active_num_decodes],
604 + cache_lens=self.max_model_len_tensor[:active_num_decodes],
605 + )
606 fp8_metadata.decode = FP8Meta.Decode(
607 seq_lens=common_attn_metadata.seq_lens[:active_num_decodes],
608 kernel_metadata=kernel_meta,
@@ -626,19 +705,35 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
705 "the bf16 sparse path is not supported under DCP."
706 )
707
708 + self.pcp_dcp_kv_gather = False
709 + self.gathered_kv_workspace: torch.Tensor | None = None
710 if kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS:
711 # Reserve workspace during initialization
712 assert vllm_config is not None and vllm_config.model_config is not None
713 prefill_workspace_size = get_prefill_workspace_size(
714 vllm_config.model_config.max_model_len
715 )
635 - self.prefill_workspace_shape = (prefill_workspace_size, head_size)
636 - self.q_concat_buffer, self.prefill_bf16_workspace = (
637 - current_workspace_manager().get_simultaneous(
638 - (q_concat_shape, torch.bfloat16),
639 - (self.prefill_workspace_shape, torch.bfloat16),
640 - )
716 + parallel_config = vllm_config.parallel_config
717 + self.pcp_dcp_kv_gather = (
718 + parallel_config.prefill_context_parallel_size > 1
719 + and parallel_config.decode_context_parallel_size > 1
720 )
721 + shard_rows = prefill_workspace_size
722 + if self.pcp_dcp_kv_gather:
723 + # PCP+DCP upconverts this rank's KV shard, then all-gathers the
724 + # shards into a workspace of the full prefill size.
725 + shard_rows //= parallel_config.decode_context_parallel_size
726 + self.prefill_workspace_shape = (shard_rows, head_size)
727 + specs = [
728 + (q_concat_shape, torch.bfloat16),
729 + (self.prefill_workspace_shape, torch.bfloat16),
730 + ]
731 + if self.pcp_dcp_kv_gather:
732 + specs.append(((prefill_workspace_size, head_size), torch.bfloat16))
733 + buffers = current_workspace_manager().get_simultaneous(*specs)
734 + self.q_concat_buffer, self.prefill_bf16_workspace = buffers[:2]
735 + if self.pcp_dcp_kv_gather:
736 + self.gathered_kv_workspace = buffers[2]
737 else:
738 (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous(
739 (q_concat_shape, torch.bfloat16),
@@ -729,13 +824,49 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
824 return attn_out, lse
825 return torch.cat([decode_out, attn_out], dim=0), None
826
827 + def _gather_prefill_chunk(
828 + self,
829 + chunk: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill.Chunk",
830 + shard: torch.Tensor,
831 + topk_indices: torch.Tensor,
832 + attn_metadata: FlashMLASparseMetadata,
833 + prefill_meta: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill",
834 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
835 + """All-gather this rank's upconverted KV shard so the chunk's rows attend
836 + the whole context, and map their top-k onto the rank-major result."""
837 + shard_rows = int(chunk.chunk_tot_seqlen)
838 + assert self.gathered_kv_workspace is not None
839 + gathered_kv = self.gathered_kv_workspace[: self.dcp_world_size * shard_rows]
840 + dist.all_gather_into_tensor(
841 + gathered_kv, shard, group=get_dcp_group().device_group
842 + )
843 +
844 + tokens = chunk.tokens_slice
845 + topk_indices, topk_length = triton_convert_req_index_to_global_index(
846 + attn_metadata.req_id_per_token[tokens],
847 + attn_metadata.block_table,
848 + topk_indices,
849 + BLOCK_SIZE=attn_metadata.block_size,
850 + NUM_TOPK_TOKENS=topk_indices.shape[1],
851 + HAS_PREFILL_WORKSPACE=True,
852 + prefill_workspace_request_ids=prefill_meta.request_ids[tokens],
853 + prefill_workspace_starts=prefill_meta.workspace_starts,
854 + prefill_workspace_rank_stride=shard_rows,
855 + dcp_size=self.dcp_world_size,
856 + dcp_rank=self.dcp_rank,
857 + cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size,
858 + return_valid_counts=True,
859 + )
860 + return gathered_kv, topk_indices, topk_length
861 +
862 def _forward_fp8_kv_separate_prefill_decode(
863 self,
864 q: torch.Tensor,
865 kv_c_and_k_pe_cache: torch.Tensor,
866 topk_indices: torch.Tensor,
867 attn_metadata: FlashMLASparseMetadata,
738 - ) -> torch.Tensor:
868 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
869 + """The lse covers the decode rows and is only returned under PCP+DCP."""
870 fp8_metadata = attn_metadata.fp8_extra_metadata
871 assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode)
872 num_decodes = fp8_metadata.num_decodes
@@ -783,8 +914,13 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
914 # For BF16 cache: always use global cache slots (no workspace)
915 # prefill_workspace_starts has been adjusted in-place per chunk so
916 # prefill indices automatically come out chunk-local
786 - topk_length = None
787 - if num_prefill_tokens == 0 and not uses_host_cache:
917 + # Under PCP+DCP the decode rows are converted by the mixed-batch call
918 + # and the prefill rows per chunk, against the gathered workspace.
919 + lse: torch.Tensor | None = None
920 + topk_length: torch.Tensor | None = None
921 + if self.pcp_dcp_kv_gather:
922 + pass
923 + elif num_prefill_tokens == 0 and not uses_host_cache:
924 topk_indices, topk_length = self._convert_logical_to_physical_topk(
925 topk_indices,
926 attn_metadata,
@@ -812,6 +948,17 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
948 topk_indices: torch.Tensor,
949 ) -> torch.Tensor:
950 assert fp8_metadata.decode is not None
951 + if self.pcp_dcp_kv_gather:
952 + # Attend this rank's KV shard; the LSE feeds the DCP merge.
953 + nonlocal lse
954 + out, lse = self._forward_fp8_kv_mixed_batch(
955 + q,
956 + kv_c_and_k_pe_cache,
957 + topk_indices,
958 + attn_metadata,
959 + kernel_metadata=fp8_metadata.decode.kernel_metadata,
960 + )
961 + return out
962 if uses_host_cache:
963 return self._host_backed_fp8_decode(
964 q,
@@ -896,9 +1043,20 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1043 )
1044
1045 chunk_q = q[chunk.tokens_slice]
899 - chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice]
900 - assert topk_length is not None
901 - chunk_topk_length = topk_length[chunk.tokens_slice]
1046 + if self.pcp_dcp_kv_gather:
1047 + chunk_workspace, chunk_topk_indices_workspace, chunk_topk_length = (
1048 + self._gather_prefill_chunk(
1049 + chunk,
1050 + chunk_workspace,
1051 + topk_indices[chunk.tokens_slice],
1052 + attn_metadata,
1053 + fp8_metadata.prefill,
1054 + )
1055 + )
1056 + else:
1057 + chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice]
1058 + assert topk_length is not None
1059 + chunk_topk_length = topk_length[chunk.tokens_slice]
1060
1061 attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel(
1062 chunk_q,
@@ -907,7 +1065,10 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1065 chunk_topk_length,
1066 )
1067
910 - return attn_out
1068 + if self.pcp_dcp_kv_gather and lse is None:
1069 + # No decode rows: the DCP merge still expects an LSE, of no rows.
1070 + lse = q.new_empty((0, q.shape[1]), dtype=torch.float32)
1071 + return attn_out, lse
1072
1073 def _forward_fp8_kv_mixed_batch(
1074 self,
@@ -915,6 +1076,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1076 kv_c_and_k_pe_cache: torch.Tensor,
1077 topk_indices: torch.Tensor,
1078 attn_metadata: FlashMLASparseMetadata,
1079 + kernel_metadata: "FlashMLASparseMetadata.FP8KernelMetadata | None" = None,
1080 ) -> tuple[torch.Tensor, torch.Tensor | None]:
1081 """Mixed batch FP8 forward path that treats all tokens as one batch.
1082
@@ -924,12 +1086,8 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1086
1087 The lse is only returned when DCP needs it, otherwise None.
1088 """
927 - assert attn_metadata.fp8_extra_metadata is not None
928 - assert isinstance(
929 - attn_metadata.fp8_extra_metadata,
930 - FlashMLASparseMetadata.FP8KernelMetadata,
931 - )
932 - fp8_metadata = attn_metadata.fp8_extra_metadata
1089 + fp8_metadata = kernel_metadata or attn_metadata.fp8_extra_metadata
1090 + assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8KernelMetadata)
1091
1092 block_table = attn_metadata.block_table
1093 # req_id_per_token covers the whole batch; slice it to the MQA tokens
@@ -1160,7 +1318,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1318 q, kv_c_and_k_pe_cache, topk_indices, attn_metadata
1319 )
1320 else:
1163 - attn_out = self._forward_fp8_kv_separate_prefill_decode(
1321 + attn_out, lse = self._forward_fp8_kv_separate_prefill_decode(
1322 q, kv_c_and_k_pe_cache, topk_indices, attn_metadata
1323 )
1324
vllm/v1/attention/backends/mla/indexer.py
+210 -20
@@ -3,6 +3,7 @@
3 from dataclasses import dataclass
4 from typing import Any
5
6 +import numpy as np
7 import torch
8
9 import vllm.envs as envs
@@ -25,6 +26,7 @@ from vllm.utils.deep_gemm import (
26 native_next_n_supported,
27 )
28 from vllm.utils.platform_utils import num_compute_units
29 +from vllm.utils.torch_utils import PIN_MEMORY
30 from vllm.v1.attention.backend import (
31 AttentionBackend,
32 AttentionCGSupport,
@@ -33,6 +35,7 @@ from vllm.v1.attention.backend import (
35 MultipleOf,
36 )
37 from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
38 +from vllm.v1.attention.backends.mla.sparse_utils import request_row_bounds
39 from vllm.v1.attention.backends.utils import (
40 get_dcp_local_seq_lens,
41 split_decodes_and_prefills,
@@ -43,6 +46,7 @@ from vllm.v1.kv_cache_interface import (
46 KVCacheSpec,
47 MLAAttentionSpec,
48 )
49 +from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
50
51 logger = init_logger(__name__)
52
@@ -259,6 +263,87 @@ class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend):
263 return [64 if current_platform.is_device_capability_family(90) else 128]
264
265
266 +@dataclass(frozen=True)
267 +class PCPGlobalChunkPlan:
268 + """PCP packing for one indexer prefill chunk under PCP + DCP."""
269 +
270 + # [num_reqs+1] cumsum of the padded per-request context (shard rows x W):
271 + # the global layout.
272 + row_start_cu: torch.Tensor
273 + # [num_reqs+1] like row_start_cu, but only a region's FIRST row carries its
274 + # extent.
275 + global_cu: torch.Tensor
276 + # [num_reqs+1] cumsum of shard rows: this rank's padded layout.
277 + padded_local_cu: torch.Tensor
278 + padded_local_total: int
279 + total: int
280 + # [total] for each global row, its position in the rank-major gathered buffer.
281 + deinterleave_idx: torch.Tensor
282 +
283 +
284 +def build_pcp_global_chunk_plan(
285 + row_req_idx: np.ndarray,
286 + row_shard_rows: np.ndarray,
287 + dcp_world_size: int,
288 + device: torch.device,
289 +) -> PCPGlobalChunkPlan:
290 + """Plan the PCP packing for one chunk from its requests' KV shard rows.
291 +
292 + The global layout is padded to ``shard rows x W`` per request. Positions
293 + past a request's true extent lie beyond every token's causal bound, so
294 + the kernel never reads them.
295 + """
296 + shard_rows = np.ascontiguousarray(row_shard_rows, dtype=np.int64)
297 + assert row_req_idx.shape == shard_rows.shape
298 + num_rows = len(shard_rows)
299 +
300 + row_bounds = request_row_bounds(row_req_idx)
301 + region_first_row = row_bounds[:-1]
302 + region_of_row = np.repeat(np.arange(len(region_first_row)), np.diff(row_bounds))
303 +
304 + region_padded = shard_rows[region_first_row]
305 + assert np.all(region_padded > 0), (
306 + f"PCP+DCP prefill got an empty context: {region_padded.tolist()}"
307 + )
308 +
309 + # Pinned host buffers, filled in place so the copies below are async.
310 + cu = torch.zeros((3, num_rows + 1), dtype=torch.int32, pin_memory=PIN_MEMORY)
311 + row_start_cu, global_cu, padded_cu = cu
312 + # Only a region's first row carries its extent.
313 + first_row = torch.from_numpy(region_first_row)
314 + row_padded = torch.zeros(num_rows, dtype=torch.int32)
315 + row_padded[first_row] = torch.from_numpy(region_padded).int()
316 + torch.cumsum(row_padded * dcp_world_size, 0, out=global_cu[1:])
317 + torch.cumsum(row_padded, 0, out=padded_cu[1:])
318 + row_start_cu[:num_rows] = global_cu[first_row][torch.from_numpy(region_of_row)]
319 + row_start_cu[num_rows] = global_cu[num_rows]
320 + total = int(global_cu[num_rows])
321 + padded_total = int(padded_cu[num_rows])
322 +
323 + idx = torch.empty(total, dtype=torch.int64, pin_memory=PIN_MEMORY)
324 + idx_np = idx.numpy()
325 + region_start = global_cu[first_row].numpy()
326 + region_padded_start = padded_cu[first_row].numpy()
327 + for i in range(len(region_first_row)):
328 + g = int(region_padded[i]) * dcp_world_size
329 + t = np.arange(g, dtype=np.int64)
330 + idx_np[region_start[i] : region_start[i] + g] = (
331 + (t % dcp_world_size) * padded_total
332 + + region_padded_start[i]
333 + + t // dcp_world_size
334 + )
335 +
336 + cu = async_copy_to_gpu(cu, device=device)
337 + return PCPGlobalChunkPlan(
338 + row_start_cu=cu[0],
339 + global_cu=cu[1],
340 + padded_local_cu=cu[2],
341 + padded_local_total=padded_total,
342 + total=total,
343 + deinterleave_idx=async_copy_to_gpu(idx, device=device),
344 + )
345 +
346 +
347 @dataclass
348 class DeepseekV32IndexerPrefillChunkMetadata:
349 block_table: torch.Tensor
@@ -277,6 +362,8 @@ class DeepseekV32IndexerPrefillChunkMetadata:
362 local_total_seq_lens: int = 0
363 max_local_total_seq_lens: int = 0
364
365 + pcp_deinterleave_idx: torch.Tensor | None = None
366 +
367
368 class BuildPrefillChunkMetadataKernel(
369 VllmTritonJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]
@@ -960,6 +1047,48 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1047 self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0
1048 return self.global_decode_seq_lens_buffer[:num_decode_tokens]
1049
1050 + def _split_pcp_dcp_prefill_chunks(
1051 + self,
1052 + row_req_idx: np.ndarray,
1053 + row_shard_rows: np.ndarray,
1054 + row_query_lens_cpu: torch.Tensor,
1055 + max_logits_bytes: int,
1056 + request_offset: int,
1057 + ) -> list[tuple[slice, slice]]:
1058 + """Chunk by request rather than by row, so a split prefill's two rows
1059 + charge their shared context once, then widen each chunk back to rows.
1060 +
1061 + ``row_shard_rows`` holds the whole request's largest DCP shard on
1062 + every row, so the plan is identical on every PCP rank.
1063 + """
1064 + row_bounds = request_row_bounds(row_req_idx)
1065 + first_rows = row_bounds[:-1]
1066 + # Each request's context, padded to whole DCP shards.
1067 + seq_lens = row_shard_rows[first_rows] * self.dcp_world_size
1068 + # Every rank holds a full-size chunk of a split request (shorter ones
1069 + # are replicated), so rows x longest row is the same on every rank,
1070 + # whichever row holds the short tail.
1071 + row_query_lens = row_query_lens_cpu.numpy()
1072 + query_lens = np.diff(row_bounds) * np.maximum.reduceat(
1073 + row_query_lens, first_rows
1074 + )
1075 + chunk_specs = self._split_indexer_prefill_chunks(
1076 + torch.from_numpy(seq_lens.astype(np.int32)),
1077 + torch.from_numpy(query_lens.astype(np.int32)),
1078 + self.max_prefill_buffer_size,
1079 + max_logits_bytes,
1080 + )
1081 + return [
1082 + (
1083 + slice(
1084 + request_offset + int(row_bounds[request_slice.start]),
1085 + request_offset + int(row_bounds[request_slice.stop]),
1086 + ),
1087 + query_slice,
1088 + )
1089 + for request_slice, query_slice in chunk_specs
1090 + ]
1091 +
1092 @staticmethod
1093 def _split_indexer_prefill_chunks(
1094 compressed_seq_lens_cpu: torch.Tensor,
@@ -1094,16 +1223,52 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1223 # slice below).
1224 assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
1225 seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
1097 - chunk_specs = self._split_indexer_prefill_chunks(
1098 - compressed_seq_lens_cpu[num_decodes:],
1099 - prefill_query_lens_cpu,
1100 - self.max_prefill_buffer_size,
1101 - max_logits_bytes,
1102 - request_offset=num_decodes,
1103 - )
1226 + req_idx = None
1227 + shard_rows = None
1228 + if self.use_pcp and self.dcp_world_size > 1:
1229 + # The gathered KV must be packed identically on every PCP rank:
1230 + # chunk by request from its DCP shard rows, which every rank
1231 + # holds. A dummy batch bypasses the PCP manager and has one
1232 + # row per request, so its own extent is the request's.
1233 + req_idx = common_attn_metadata.req_idx
1234 + if req_idx is None:
1235 + req_idx = np.arange(num_reqs)
1236 + shard_rows_cpu = common_attn_metadata.dcp_local_seq_lens_cpu_upper_bound
1237 + if shard_rows_cpu is None:
1238 + shard_rows_cpu = get_dcp_local_seq_lens(
1239 + seq_lens_cpu,
1240 + self.dcp_world_size,
1241 + 0,
1242 + self.cp_kv_cache_interleave_size,
1243 + )
1244 + shard_rows = shard_rows_cpu.numpy()
1245 + chunk_specs = self._split_pcp_dcp_prefill_chunks(
1246 + req_idx[num_decodes:],
1247 + shard_rows[num_decodes:],
1248 + prefill_query_lens_cpu,
1249 + max_logits_bytes,
1250 + request_offset=num_decodes,
1251 + )
1252 + else:
1253 + chunk_specs = self._split_indexer_prefill_chunks(
1254 + compressed_seq_lens_cpu[num_decodes:],
1255 + prefill_query_lens_cpu,
1256 + self.max_prefill_buffer_size,
1257 + max_logits_bytes,
1258 + request_offset=num_decodes,
1259 + )
1260
1261 chunks = []
1262 for req_slice, query_slice in chunk_specs:
1263 + pcp_plan = None
1264 + if req_idx is not None:
1265 + assert shard_rows is not None
1266 + pcp_plan = build_pcp_global_chunk_plan(
1267 + req_idx[req_slice],
1268 + shard_rows[req_slice],
1269 + self.dcp_world_size,
1270 + self.device,
1271 + )
1272 metadata = build_prefill_chunk_metadata(
1273 req_slice.start,
1274 req_slice.stop,
@@ -1119,6 +1284,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1284 dcp_rank=self.dcp_rank,
1285 dcp_world_size=self.dcp_world_size,
1286 cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
1287 + pcp_plan=pcp_plan,
1288 )
1289 # Skip when total_seq_lens is 0 (i.e., no compressed token).
1290 if metadata is not None:
@@ -1353,8 +1519,12 @@ def build_prefill_chunk_metadata(
1519 dcp_rank: int = 0,
1520 dcp_world_size: int = 1,
1521 cp_kv_cache_interleave_size: int = 1,
1522 + pcp_plan: PCPGlobalChunkPlan | None = None,
1523 ) -> DeepseekV32IndexerPrefillChunkMetadata | None:
1357 - total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item()
1524 + if pcp_plan is not None:
1525 + total_seq_lens = pcp_plan.total
1526 + else:
1527 + total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item()
1528 if total_seq_lens == 0:
1529 return None
1530
@@ -1362,15 +1532,22 @@ def build_prefill_chunk_metadata(
1532 device = block_table.device
1533 token_to_seq = torch.empty(total_seq_lens, dtype=torch.int32, device=device)
1534
1365 - cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device)
1366 - # Assigning to slice avoids cpu sync.
1367 - cu_seq_lens[:1] = 0
1368 - torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:])
1535 + if pcp_plan is not None:
1536 + cu_seq_lens = pcp_plan.global_cu
1537 + else:
1538 + cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device)
1539 + # Assigning to slice avoids cpu sync.
1540 + cu_seq_lens[:1] = 0
1541 + torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:])
1542
1543 local_cu_seq_lens = cu_seq_lens
1544 local_total_seq_lens = total_seq_lens
1545 max_local_total_seq_lens = total_seq_lens
1373 - if dcp_world_size > 1:
1546 + if pcp_plan is not None:
1547 + local_cu_seq_lens = pcp_plan.padded_local_cu
1548 + local_total_seq_lens = pcp_plan.padded_local_total
1549 + max_local_total_seq_lens = pcp_plan.padded_local_total
1550 + elif dcp_world_size > 1:
1551 # Per-rank local KV length under interleave-aware DCP sharding, shape
1552 # [num_reqs, dcp_world_size]. Reuse the canonical CP helper so the
1553 # sharding matches the rest of the DCP pipeline (decode/prefill).
@@ -1394,35 +1571,47 @@ def build_prefill_chunk_metadata(
1571 (query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item()
1572 )
1573 if query_slice is not None:
1397 - qs_start = query_slice.start
1398 - qs_stop = query_slice.stop
1574 + qs_start = min(query_slice.start, total_query_len)
1575 + qs_stop = min(query_slice.stop, total_query_len)
1576 else:
1577 qs_start = 0
1578 qs_stop = total_query_len
1579 output_query_len = qs_stop - qs_start
1580 + assert 0 <= output_query_len <= total_query_len - qs_start, (
1581 + f"query slice [{qs_start}, {qs_stop}) does not fit this rank's "
1582 + f"{total_query_len} tokens; the metadata tail would stay uninitialized"
1583 + )
1584
1585 cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
1586 cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)
1587
1407 - # Under DCP the kernel writes this rank's local row bounds into
1408 - # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens.
1588 + if pcp_plan is not None:
1589 + row_start_cu = pcp_plan.row_start_cu
1590 + kernel_dcp_rank, kernel_dcp_world = 0, 1
1591 + else:
1592 + # Under DCP the kernel writes this rank's local row bounds into
1593 + # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens.
1594 + row_start_cu = local_cu_seq_lens
1595 + kernel_dcp_rank, kernel_dcp_world = dcp_rank, dcp_world_size
1596 _BUILD_PREFILL_CHUNK_METADATA_KERNEL(
1597 query_start_loc,
1598 uncompressed_seq_lens[start_idx:end_idx],
1599 cu_seq_lens,
1413 - local_cu_seq_lens,
1600 + row_start_cu,
1601 token_to_seq,
1602 cu_seq_len_ks,
1603 cu_seq_len_ke,
1604 qs_start,
1605 qs_stop,
1419 - dcp_rank,
1420 - dcp_world_size,
1606 + kernel_dcp_rank,
1607 + kernel_dcp_world,
1608 cp_kv_cache_interleave_size,
1609 num_reqs=num_reqs,
1610 COMPRESS_RATIO=compress_ratio,
1611 )
1612
1613 + pcp_deinterleave_idx = pcp_plan.deinterleave_idx if pcp_plan is not None else None
1614 +
1615 token_start = query_start_loc_cpu[start_idx].item()
1616 if query_slice is not None:
1617 token_end = token_start + qs_stop
@@ -1432,6 +1621,7 @@ def build_prefill_chunk_metadata(
1621 token_end = query_start_loc_cpu[end_idx].item()
1622
1623 return DeepseekV32IndexerPrefillChunkMetadata(
1624 + pcp_deinterleave_idx=pcp_deinterleave_idx,
1625 cu_seqlen_ks=cu_seq_len_ks,
1626 cu_seqlen_ke=cu_seq_len_ke,
1627 cu_seq_lens=cu_seq_lens,
vllm/v1/attention/backends/mla/sparse_utils.py
+52 -7
@@ -5,6 +5,7 @@
5 from dataclasses import dataclass
6 from typing import Any
7
8 +import numpy as np
9 import torch
10
11 from vllm.config import VllmConfig
@@ -24,6 +25,21 @@ from vllm.utils.math_utils import cdiv
25 from vllm.v1.worker.block_table import get_block_table_width
26
27
28 +def request_row_bounds(req_idx: np.ndarray) -> np.ndarray:
29 + """Bounds of the runs of adjacent rows that belong to one request: run
30 + ``r`` is rows ``[bounds[r], bounds[r + 1])``.
31 +
32 + Under PCP a rank holds two adjacent chunk rows of a split prefill; the
33 + sparse backends give such a run one KV region.
34 + """
35 + assert req_idx.size > 0
36 + bounds = np.flatnonzero(np.r_[True, req_idx[1:] != req_idx[:-1], True])
37 + assert bounds.size - 1 == np.unique(req_idx).size, (
38 + "rows of one request must be adjacent"
39 + )
40 + return bounds
41 +
42 +
43 def flat_kv_row_view(
44 kv_cache: torch.Tensor, # [num_blocks, block_size, head_dim]
45 block_size: int,
@@ -69,7 +85,7 @@ class ConvertReqIndexToGlobalIndexKernel(
85 block_table_stride: int
86
87 @staticmethod
72 - @triton.jit(do_not_specialize=["max_num_blocks_per_req"])
88 + @triton.jit(do_not_specialize=["max_num_blocks_per_req", "workspace_rank_stride"])
89 def kernel(
90 req_id_ptr, # int32 [num_tokens]
91 block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]
@@ -78,6 +94,7 @@ class ConvertReqIndexToGlobalIndexKernel(
94 valid_count_ptr, # int32 [num_tokens] - output valid count per row
95 prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill
96 workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr
97 + workspace_rank_stride, # rows per DCP rank of a gathered prefill workspace
98 # shapes (compile-time where possible)
99 max_num_blocks_per_req,
100 BLOCK_SIZE: tl.constexpr,
@@ -146,7 +163,7 @@ class ConvertReqIndexToGlobalIndexKernel(
163 # Guard block_table access
164 valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)
165 bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
149 - is_invalid_tok |= ~valid_block | is_remote
166 + is_invalid_tok |= ~valid_block | (is_remote & ~is_prefill)
167 base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0)
168 out_val = base * BLOCK_STRIDE_ROWS + inblock_off
169
@@ -155,7 +172,11 @@ class ConvertReqIndexToGlobalIndexKernel(
172 workspace_start = tl.load(
173 workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0
174 )
158 - prefill_out = workspace_start + tok
175 + # Under DCP the prefill workspace is the all-gather of every rank's
176 + # shard, rank-major. With DCP_SIZE == 1 this is workspace_start + tok.
177 + prefill_out = (
178 + owning_rank * workspace_rank_stride + workspace_start + local_idx
179 + )
180 out_val = tl.where(is_prefill, prefill_out, out_val)
181 out_val = tl.where(is_invalid_tok, -1, out_val)
182
@@ -238,6 +259,22 @@ class ConvertReqIndexToGlobalIndexKernel(
259 block_size * dcp_size,
260 )
261 max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size)
262 + pcp_dcp_combos = []
263 + if (
264 + vllm_config.parallel_config.prefill_context_parallel_size > 1
265 + and dcp_size > 1
266 + ):
267 + # PCP + DCP attends prefill rows over a DCP-gathered workspace.
268 + pcp_dcp_combos.append(
269 + dict(
270 + HAS_PREFILL_WORKSPACE=True,
271 + COUNT_VALID=True,
272 + COMPACT_TO_FRONT=True,
273 + DCP_SIZE=dcp_size,
274 + DCP_RANK=dcp_rank,
275 + DCP_INTERLEAVE=dcp_interleave,
276 + )
277 + )
278 return self._trace_dispatch(self.dispatch)(
279 zip_inputs(
280 dict(
@@ -280,6 +317,7 @@ class ConvertReqIndexToGlobalIndexKernel(
317 DCP_RANK=dcp_rank,
318 DCP_INTERLEAVE=dcp_interleave,
319 ),
320 + *pcp_dcp_combos,
321 ),
322 BLOCK_SIZE=block_size,
323 BLOCK_STRIDE_ROWS=block_stride_rows,
@@ -316,6 +354,7 @@ class ConvertReqIndexToGlobalIndexKernel(
354 int32_ptr if compile_key.has_prefill_workspace else None
355 ),
356 max_num_blocks_per_req=1,
357 + workspace_rank_stride=1,
358 BLOCK_SIZE=compile_key.block_size,
359 BLOCK_STRIDE_ROWS=compile_key.block_stride_rows,
360 BLOCK_N=compile_key.block_n,
@@ -339,6 +378,7 @@ class ConvertReqIndexToGlobalIndexKernel(
378 prefill_workspace_request_ids: torch.Tensor | None,
379 prefill_workspace_starts: torch.Tensor | None,
380 *,
381 + workspace_rank_stride: int,
382 max_num_blocks_per_req: int,
383 BLOCK_SIZE: int,
384 BLOCK_STRIDE_ROWS: int,
@@ -406,6 +446,10 @@ def triton_convert_req_index_to_global_index(
446 HAS_PREFILL_WORKSPACE: bool = False,
447 prefill_workspace_request_ids: torch.Tensor | None = None,
448 prefill_workspace_starts: torch.Tensor | None = None,
449 + prefill_workspace_rank_stride: int | None = None,
450 + dcp_size: int = 1,
451 + dcp_rank: int = 0,
452 + cp_kv_cache_interleave_size: int = 1,
453 return_valid_counts: bool = False,
454 out: torch.Tensor | None = None,
455 valid_counts_out: torch.Tensor | None = None,
@@ -511,6 +555,7 @@ def triton_convert_req_index_to_global_index(
555 prefill_workspace_request_ids,
556 prefill_workspace_starts,
557 # shapes / constexprs
558 + workspace_rank_stride=prefill_workspace_rank_stride or 0,
559 max_num_blocks_per_req=max_num_blocks_per_req,
560 BLOCK_SIZE=BLOCK_SIZE,
561 BLOCK_STRIDE_ROWS=(
@@ -521,10 +566,9 @@ def triton_convert_req_index_to_global_index(
566 HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE,
567 COUNT_VALID=return_valid_counts,
568 COMPACT_TO_FRONT=return_valid_counts,
524 - # DCP disabled (no-op de-interleave)
525 - DCP_SIZE=1,
526 - DCP_RANK=0,
527 - DCP_INTERLEAVE=1,
569 + DCP_SIZE=dcp_size,
570 + DCP_RANK=dcp_rank,
571 + DCP_INTERLEAVE=cp_kv_cache_interleave_size,
572 )
573
574 if return_valid_counts:
@@ -618,6 +662,7 @@ def triton_filter_and_convert_dcp_index(
662 # No prefill workspace on the DCP decode path.
663 None,
664 None,
665 + workspace_rank_stride=0,
666 max_num_blocks_per_req=max_num_blocks_per_req,
667 BLOCK_SIZE=BLOCK_SIZE,
668 BLOCK_STRIDE_ROWS=(
vllm/v1/worker/gpu/attn_utils.py
+8
@@ -400,6 +400,7 @@ def build_attn_metadata(
400 kv_cache_config: KVCacheConfig,
401 seq_lens_cpu_upper_bound: torch.Tensor | None = None,
402 dcp_local_seq_lens: torch.Tensor | None = None,
403 + dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None,
404 positions: torch.Tensor | None = None,
405 is_prefilling: torch.Tensor | None = None,
406 mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None,
@@ -409,10 +410,15 @@ def build_attn_metadata(
410 rswa_prefix_lens: torch.Tensor | None = None,
411 ubatch_idx: int = 0,
412 fast_prefill: FastPrefillBatchMetadata | None = None,
413 + req_idx: np.ndarray | None = None,
414 ) -> dict[str, Any]:
415 seq_lens = seq_lens[:num_reqs]
416 if dcp_local_seq_lens is not None:
417 dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs]
418 + if dcp_local_seq_lens_cpu_upper_bound is not None:
419 + dcp_local_seq_lens_cpu_upper_bound = dcp_local_seq_lens_cpu_upper_bound[
420 + :num_reqs
421 + ]
422 if seq_lens_cpu_upper_bound is not None:
423 seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound[:num_reqs]
424
@@ -457,10 +463,12 @@ def build_attn_metadata(
463 slot_mapping=slot_mapping,
464 causal=group_causal,
465 dcp_local_seq_lens=dcp_local_seq_lens,
466 + dcp_local_seq_lens_cpu_upper_bound=dcp_local_seq_lens_cpu_upper_bound,
467 positions=positions,
468 is_prefilling=group_is_prefilling,
469 mm_req_doc_ranges=mm_req_doc_ranges,
470 rswa_prefix_lens=rswa_prefix_lens,
471 + req_idx=req_idx,
472 **common_attn_metadata_extra_kwargs,
473 )
474
vllm/v1/worker/gpu/input_batch.py
+3
@@ -114,6 +114,9 @@ class InputBatch:
114 # (cudagraph capture) batches, which run the KV-sharing layers in full.
115 fast_prefill: "FastPrefillBatchMetadata | None" = None
116
117 + # [num_reqs] set only under PCP+DCP (see CommonAttentionMetadata).
118 + dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None
119 +
120 @classmethod
121 def make_dummy(
122 cls,
vllm/v1/worker/gpu/model_states/default.py
+2
@@ -220,11 +220,13 @@ class DefaultModelState(ModelState):
220 kv_cache_config=kv_cache_config,
221 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
222 dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
223 + dcp_local_seq_lens_cpu_upper_bound=input_batch.dcp_local_seq_lens_cpu_upper_bound,
224 positions=input_batch.positions,
225 is_prefilling=torch.from_numpy(input_batch.is_prefilling_np),
226 mm_req_doc_ranges=req_doc_ranges,
227 for_cudagraph_capture=for_capture,
228 rswa_prefix_lens=input_batch.prompt_lens,
229 + req_idx=input_batch.idx_mapping_np,
230 ubatch_idx=ubatch_idx,
231 fast_prefill=input_batch.fast_prefill,
232 )
vllm/v1/worker/gpu/pcp_manager.py
+70 -26
@@ -6,10 +6,10 @@ from dataclasses import dataclass, replace
6 import numpy as np
7 import torch
8
9 -from vllm.config import VllmConfig
9 +from vllm.config import CUDAGraphMode, VllmConfig
10 from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group
11 from vllm.logger import init_logger
12 -from vllm.v1.attention.backends.utils import PAD_SLOT_ID
12 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens
13 from vllm.v1.worker.gpu.block_table import BlockTables
14 from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
15 from vllm.v1.worker.gpu.input_batch import (
@@ -154,30 +154,48 @@ class PCPManager:
154 "MRV2 PCP only supports DSpark or single-module MTP "
155 "speculative decoding."
156 )
157 - if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs():
157 + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode
158 + is_sparse_mla = hasattr(model_config.hf_text_config, "index_topk")
159 + if parallel_config.decode_context_parallel_size > 1 and not is_sparse_mla:
160 + # Dense MLA prefill sizes its DCP KV gather from each rank's own
161 + # chunk rows, so the ranks' collectives diverge (#53573).
162 + raise NotImplementedError("MRV2 PCP + DCP supports sparse MLA models only.")
163 + if (
164 + is_sparse_mla
165 + and parallel_config.decode_context_parallel_size == 1
166 + and cudagraph_mode != CUDAGraphMode.NONE
167 + ):
168 + raise NotImplementedError(
169 + "MRV2 sparse MLA PCP does not support CUDA graphs yet. "
170 + "Set -cc.cudagraph_mode=NONE."
171 + )
172 + if cudagraph_mode.has_full_cudagraphs():
173 raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.")
174 + if (
175 + parallel_config.decode_context_parallel_size > 1
176 + and parallel_config.dcp_comm_backend != "ag_rs"
177 + ):
178 + raise NotImplementedError(
179 + "MRV2 PCP + DCP requires dcp_comm_backend='ag_rs'; got "
180 + f"'{parallel_config.dcp_comm_backend}'."
181 + )
182
183 @staticmethod
184 def _reorder_segments(
185 segments: list[RankSegment],
163 - num_computed_tokens: np.ndarray,
186 is_prefilling: np.ndarray,
165 - query_start_loc_np: np.ndarray,
187 ) -> list[RankSegment]:
167 - """Move pure prefills last to match the batch ordering expected by
168 - attention backends like MLA and sparse MLA.
169 - """
188 + """Order this rank's rows decodes-first, then prefills, canonically."""
189
171 - def is_pure_prefill(segment: RankSegment) -> bool:
190 + def sort_key(segment: RankSegment) -> tuple[bool, int, int]:
191 req_idx = segment.global_batch_req_idx
173 - start_pos = (
174 - num_computed_tokens[req_idx]
175 - + segment.global_batch_slice.start
176 - - query_start_loc_np[req_idx]
192 + return (
193 + bool(is_prefilling[req_idx]),
194 + req_idx,
195 + segment.global_batch_slice.start,
196 )
178 - return is_prefilling[req_idx] and start_pos == 0
197
180 - segments.sort(key=is_pure_prefill)
198 + segments.sort(key=sort_key)
199 rank_offset = 0
200 for index, segment in enumerate(segments):
201 segments[index] = replace(
@@ -189,6 +207,21 @@ class PCPManager:
207 rank_offset += segment.num_tokens
208 return segments
209
210 + def replicated_requests(
211 + self,
212 + num_scheduled_tokens: np.ndarray,
213 + is_prefilling: np.ndarray,
214 + ) -> np.ndarray:
215 + """Per global request, whether every PCP rank gets the whole query."""
216 + num_chunks = 2 * self.pcp_world_size
217 + query_lens = np.asarray(num_scheduled_tokens, dtype=np.int64)
218 + replicated = ~np.asarray(is_prefilling, dtype=np.bool_)
219 + if self.dcp_world_size > 1:
220 + chunk_sizes = (query_lens + num_chunks - 1) // num_chunks
221 + drops_a_chunk = (num_chunks - 1) * chunk_sizes >= query_lens
222 + replicated |= drops_a_chunk
223 + return replicated
224 +
225 def _iter_rank_chunks(
226 self,
227 rank: int,
@@ -204,17 +237,20 @@ class PCPManager:
237 rank 1: 1 6
238 rank 2: 2 5
239 rank 3: 3 4
240 +
241 + Decodes, and prefills too short to fill all eight, are replicated instead.
242 """
243 num_chunks = 2 * self.pcp_world_size
244 + replicated = self.replicated_requests(num_scheduled_tokens, is_prefilling)
245 for global_batch_req_idx, num_tokens in enumerate(num_scheduled_tokens):
246 query_len = int(num_tokens)
247 if query_len == 0:
248 continue
249 chunk_indices: tuple[int, ...]
214 - if bool(is_prefilling[global_batch_req_idx]):
250 + if not replicated[global_batch_req_idx]:
251 chunk_size = (query_len + num_chunks - 1) // num_chunks
252 chunk_indices = (rank, num_chunks - 1 - rank)
217 - else: # decodes are replicated
253 + else: # decodes, and short prefills under DCP, are replicated
254 chunk_size = query_len
255 chunk_indices = (0,)
256
@@ -229,7 +265,6 @@ class PCPManager:
265 self,
266 rank: int,
267 num_scheduled_tokens: np.ndarray,
232 - num_computed_tokens: np.ndarray,
268 is_prefilling: np.ndarray,
269 query_start_loc_np: np.ndarray,
270 ) -> list[RankSegment]:
@@ -248,12 +283,7 @@ class PCPManager:
283 )
284 )
285 rank_offset += chunk_len
251 - return self._reorder_segments(
252 - rank_segments,
253 - num_computed_tokens,
254 - is_prefilling,
255 - query_start_loc_np,
256 - )
286 + return self._reorder_segments(rank_segments, is_prefilling)
287
288 def _build_batch_layout(
289 self,
@@ -263,13 +293,13 @@ class PCPManager:
293 query_start_loc_np: np.ndarray,
294 padded_num_tokens: int | None = None,
295 ) -> tuple[list[list[RankSegment]], list[int]]:
296 + replicated = self.replicated_requests(num_scheduled_tokens, is_prefilling)
297 segments_by_rank = []
298 per_rank_num_tokens = []
299 for rank in range(self.pcp_world_size):
300 segments = self._get_rank_segments(
301 rank,
302 num_scheduled_tokens,
272 - num_computed_tokens,
303 is_prefilling,
304 query_start_loc_np,
305 )
@@ -309,7 +339,7 @@ class PCPManager:
339 dtype=np.int64,
340 )
341 # Cache insertion pairs one slot entry with each rank's local decode.
312 - if not bool(is_prefilling[segment.global_batch_req_idx]) and rank != 0:
342 + if replicated[segment.global_batch_req_idx] and rank != 0:
343 continue
344 gathered_kv_write_mask[padded_gathered_slice] = True
345 hidden_restore_idx[segment.global_batch_slice] = np.arange(
@@ -527,6 +557,19 @@ class PCPManager:
557 )
558 seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32)
559 seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens
560 + dcp_local_seq_lens_cpu_upper_bound = None
561 + if self.dcp_world_size > 1:
562 + # The largest DCP shard of each row's whole request, identical on
563 + # every PCP rank: the sparse backends pad their KV gather to it.
564 + request_seq_lens = (num_computed_tokens + num_scheduled_tokens)[
565 + local_to_global_batch_req_idx_np
566 + ]
567 + dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens(
568 + torch.from_numpy(request_seq_lens.astype(np.int32)),
569 + self.dcp_world_size,
570 + 0,
571 + self.cp_interleave,
572 + )
573
574 self._local_batch = replace(
575 input_batch,
@@ -549,6 +592,7 @@ class PCPManager:
592 seq_lens=seq_lens,
593 seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np),
594 dcp_local_seq_lens=None,
595 + dcp_local_seq_lens_cpu_upper_bound=dcp_local_seq_lens_cpu_upper_bound,
596 num_computed_tokens_np=local_start_pos_np,
597 prefill_len_np=local_prefill_len_np,
598 num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np,
vllm/v1/worker/gpu_model_runner.py
+1 -1
@@ -2486,7 +2486,7 @@ class GPUModelRunner(
2486 self.dcp_local_seq_lens.copy_to_gpu(num_reqs_padded)
2487
2488 cm_base.dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs_padded]
2489 - cm_base.dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[
2489 + cm_base.dcp_local_seq_lens_cpu_upper_bound = self.dcp_local_seq_lens.cpu[
2490 :num_reqs_padded
2491 ]
2492