[PCP][DCP] Enable PCP+DCP on sparse-MLA models

.buildkite/test_areas/lm_eval.yaml

@@ -110,12 +110,17 @@ steps:
   - csrc/
   - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
   - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
+  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml
   - tests/evals/gsm8k/configs/models-pcp.txt
   - vllm/model_executor/layers/quantization
   - vllm/config/parallel.py
   - vllm/distributed/parallel_state.py
   - vllm/model_executor/layers/attention/mla_attention.py
   - vllm/model_executor/layers/attention/pcp.py
+  - vllm/model_executor/layers/attention/sparse_mla_attention.py
+  - vllm/model_executor/layers/sparse_attn_indexer.py
+  - vllm/v1/attention/backends/mla/flashmla_sparse.py
+  - vllm/v1/attention/backends/mla/indexer.py
   - vllm/v1/worker/gpu/model_runner.py
   - vllm/v1/worker/gpu/pcp_manager.py
   autorun_on_main: true

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

@@ -0,0 +1,19 @@
+model_name: "nvidia/GLM-5.2-NVFP4"
+accuracy_threshold: 0.90
+num_questions: 1319
+num_fewshot: 5
+max_concurrency: 100
+server_args: >-
+  --compilation-config '{"cudagraph_mode":"PIECEWISE"}'
+  --max-model-len 4096
+  --max-num-batched-tokens 32768
+  --safetensors-load-strategy prefetch
+  --moe-backend flashinfer_cutlass
+  --prefill-context-parallel-size 4
+  --decode-context-parallel-size 4
+  --enable-expert-parallel
+  --kv-cache-dtype fp8
+env:
+  PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True"
+  VLLM_LOGGING_LEVEL: "DEBUG"
+  VLLM_USE_V2_MODEL_RUNNER: "1"

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

@@ -1,2 +1,3 @@
 GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
 GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
+GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml

tests/v1/attention/test_indexer_dcp_localize.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import numpy as np
 import pytest
 import torch
 
@@ -962,3 +963,37 @@ def test_sparse_decode_dcp_short_context_matches_non_dcp():
     dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses)
     torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5)
     torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.parametrize("dcp_world_size", [2, 4, 8])
+@pytest.mark.parametrize("req_lens", [[7], [8, 8], [1, 9], [256, 1, 2730]])
+def test_pcp_plan_deinterleave_restores_global_order(dcp_world_size, req_lens):
+    """The index gather must undo DCP sharding per request"""
+    from vllm.v1.attention.backends.mla.indexer import build_pcp_global_chunk_plan
+
+    scheduled = np.array(req_lens, dtype=np.int64)
+    shard_rows = -(-scheduled // dcp_world_size)
+    rows = np.arange(len(scheduled))
+    plan = build_pcp_global_chunk_plan(
+        rows, shard_rows, dcp_world_size, torch.device("cpu")
+    )
+
+    starts = np.concatenate([[0], np.cumsum(scheduled)])
+
+    # Build each rank's padded shard exactly as the cache gather would.
+    padded_cu = plan.padded_local_cu.tolist()
+    shards = torch.zeros(dcp_world_size, plan.padded_local_total, 1)
+    for r in range(dcp_world_size):
+        for i, g in enumerate(req_lens):
+            for t in range(r, g, dcp_world_size):
+                shards[r, padded_cu[i] + t // dcp_world_size, 0] = starts[i] + t
+
+    gathered = shards.reshape(dcp_world_size * plan.padded_local_total, 1)
+    restored = gathered[plan.deinterleave_idx]
+    # The layout is padded per request; only the real positions are read.
+    row_start = plan.row_start_cu.tolist()
+    for i, g in enumerate(req_lens):
+        torch.testing.assert_close(
+            restored[row_start[i] : row_start[i] + g, 0],
+            torch.arange(starts[i], starts[i] + g, dtype=torch.float32),
+        )

tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py

@@ -212,6 +212,7 @@ def test_index_conversion_warmup_uses_physical_block_stride():
         ),
         parallel_config=SimpleNamespace(
             decode_context_parallel_size=1,
+            prefill_context_parallel_size=1,
             cp_kv_cache_interleave_size=1,
         ),
     )

tests/v1/attention/test_sparse_mla_backends.py

@@ -1080,6 +1080,7 @@ def _convert_topk(indices, metadata, **kwargs):
         _convert_logical_to_physical_topk=_convert_topk,
         index_group=None,
         index_group_index=0,
+        pcp_dcp_kv_gather=False,
     )
 
     out, _ = FlashMLASparseImpl._forward_bf16_kv(
@@ -3257,6 +3258,7 @@ def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch):
         device=torch.device(DEVICE_TYPE),
         dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE),
         max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE),
+        pcp_dcp_kv_gather=False,
     )
     query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660])
     common_metadata = SimpleNamespace(
@@ -3363,6 +3365,7 @@ def run_kernel(**kwargs):
         index_group=None,
         index_group_index=0,
         dcp_world_size=1,
+        pcp_dcp_kv_gather=False,
         need_to_return_lse_for_decode=False,
         _fp8_flash_mla_kernel=run_kernel,
         _convert_logical_to_physical_topk=(

tests/v1/worker/test_gpu_pcp_manager.py

@@ -114,6 +114,155 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
         )
 
 
+def _rank_rows(
+    pcp_rank: int,
+    pcp_world_size: int,
+    num_scheduled_tokens: np.ndarray,
+    num_computed_tokens: np.ndarray,
+    is_prefilling: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    """One PCP+DCP rank's rows as (global request, extent)."""
+    manager = PCPManager(
+        pcp_world_size=pcp_world_size,
+        pcp_rank=pcp_rank,
+        device=torch.device("cpu"),
+        dcp_world_size=2,
+    )
+    query_start_loc_np = np.concatenate([[0], np.cumsum(num_scheduled_tokens)]).astype(
+        np.int32
+    )
+    segments = manager._get_rank_segments(
+        pcp_rank, num_scheduled_tokens, is_prefilling, query_start_loc_np
+    )
+    rows = np.array([segment.global_batch_req_idx for segment in segments])
+    extents = (num_computed_tokens + num_scheduled_tokens)[rows]
+    return rows, extents
+
+
+@pytest.mark.parametrize(
+    ("num_scheduled_tokens", "num_computed_tokens"),
+    [
+        ([32, 32], [0, 0]),
+        ([32, 32, 32, 32], [0, 0, 0, 0]),
+        ([32, 32], [100, 0]),  # one continued, one fresh
+        ([64, 32, 48], [0, 512, 0]),  # ragged lengths and mixed contexts
+        # A replicated prefill BEHIND a split one. Sorting replicated rows first
+        # put request 1 ahead of request 0.
+        ([64, 3], [0, 0]),
+        ([64, 3, 64], [0, 0, 0]),
+    ],
+)
+@pytest.mark.parametrize("pcp_world_size", [2, 4])
+def test_published_row_order_is_identical_on_every_pcp_rank(
+    num_scheduled_tokens, num_computed_tokens, pcp_world_size
+):
+    """Every rank must map its rows to the same global requests, in the same order."""
+    num_scheduled_tokens = np.array(num_scheduled_tokens, dtype=np.int32)
+    num_computed_tokens = np.array(num_computed_tokens, dtype=np.int32)
+    is_prefilling = np.ones(len(num_scheduled_tokens), dtype=np.bool_)
+
+    orders = [
+        _rank_rows(
+            rank,
+            pcp_world_size,
+            num_scheduled_tokens,
+            num_computed_tokens,
+            is_prefilling,
+        )[0]
+        for rank in range(pcp_world_size)
+    ]
+    for rank, order in enumerate(orders[1:], start=1):
+        assert np.array_equal(orders[0], order), (
+            f"rank {rank} rows map to {order.tolist()}, rank 0 to {orders[0].tolist()}"
+        )
+    # Grouped by request and ascending: what the indexer plan indexes by.
+    assert np.all(np.diff(orders[0]) >= 0)
+
+
+def test_published_row_order_puts_every_decode_before_every_prefill():
+    """split_decodes_and_prefills takes the FIRST prefilling row as the boundary."""
+    # req 0 is a continued prefill, req 1 is a decode.
+    req_idx, _ = _rank_rows(
+        pcp_rank=0,
+        pcp_world_size=2,
+        num_scheduled_tokens=np.array([32, 1], dtype=np.int32),
+        num_computed_tokens=np.array([100, 20], dtype=np.int32),
+        is_prefilling=np.array([True, False], dtype=np.bool_),
+    )
+    # The decode (request 1) must be row 0.
+    assert req_idx[0] == 1
+
+
+def test_split_prefill_rows_repeat_the_request_and_its_full_extent():
+    """Equal adjacent request indices are what let a backend share a KV region."""
+    req_idx, extents = _rank_rows(
+        pcp_rank=1,
+        pcp_world_size=2,
+        num_scheduled_tokens=np.array([64, 3], dtype=np.int32),
+        num_computed_tokens=np.array([100, 0], dtype=np.int32),
+        is_prefilling=np.ones(2, dtype=np.bool_),
+    )
+    assert req_idx.tolist() == [0, 0, 1]
+    assert extents.tolist() == [164, 164, 3]
+
+
+@pytest.mark.parametrize("query_len", [1, 2, 3, 5, 6, 9])
+def test_dcp_replicates_prefills_too_short_to_split(query_len):
+    """A prefill that cannot fill 2*pcp chunks is replicated, not split."""
+    pcp_world_size = 2
+    num_scheduled_tokens = np.array([query_len], dtype=np.int32)
+    is_prefilling = np.ones(1, dtype=np.bool_)
+
+    rows_per_rank = []
+    for rank in range(pcp_world_size):
+        manager = PCPManager(
+            pcp_world_size=pcp_world_size,
+            pcp_rank=rank,
+            device=torch.device("cpu"),
+            dcp_world_size=2,
+        )
+        assert manager.replicated_requests(num_scheduled_tokens, is_prefilling)[0]
+        rows = list(
+            manager._iter_rank_chunks(rank, num_scheduled_tokens, is_prefilling)
+        )
+        assert rows == [(0, 0, query_len)]
+        rows_per_rank.append(rows)
+    assert rows_per_rank[0] == rows_per_rank[1]
+
+
+@pytest.mark.parametrize("pcp_world_size", [2, 4, 8])
+@pytest.mark.parametrize(
+    "query_len", [16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1000, 4097]
+)
+def test_pcp_first_chunk_row_is_never_short(pcp_world_size, query_len):
+    num_scheduled_tokens = np.array([query_len], dtype=np.int32)
+    is_prefilling = np.ones(1, dtype=np.bool_)
+
+    charged_per_rank = []
+    for rank in range(pcp_world_size):
+        manager = PCPManager(
+            pcp_world_size=pcp_world_size,
+            pcp_rank=rank,
+            device=torch.device("cpu"),
+            dcp_world_size=pcp_world_size,
+        )
+        chunk_lens = [
+            chunk_len
+            for _, _, chunk_len in manager._iter_rank_chunks(
+                rank, num_scheduled_tokens, is_prefilling
+            )
+        ]
+        assert chunk_lens, f"rank {rank} got no rows for {query_len=}"
+        assert chunk_lens[0] == max(chunk_lens), (
+            f"rank {rank} emitted a short first chunk for {query_len=}: {chunk_lens}"
+        )
+        charged_per_rank.append(len(chunk_lens) * chunk_lens[0])
+
+    assert len(set(charged_per_rank)) == 1, (
+        f"ranks would chunk differently for {query_len=}: {charged_per_rank}"
+    )
+
+
 def _make_global_decode_batch(
     num_computed_tokens: list[int], buffers: InputBuffers, device: torch.device
 ) -> InputBatch:

vllm/model_executor/layers/attention/mla_attention.py

@@ -420,6 +420,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):
     """
 
     supports_dense_mha_prefill: ClassVar[bool] = True
+    # Under PCP+DCP only the decode rows carry an LSE; the base forward
+    # merges a full-batch LSE, so subclasses opt in with their own forward.
+    supports_pcp_dcp: ClassVar[bool] = False
 
     def __init__(
         self,
@@ -660,6 +663,10 @@ def __init__(
 
                     _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()
 
+        if self.use_pcp and self.impl.dcp_world_size > 1 and not self.supports_pcp_dcp:
+            raise NotImplementedError(
+                f"{type(self).__name__} does not support PCP+DCP."
+            )
         self.dcp_manager: MLADCPManager | None = None
         if self.impl.dcp_world_size > 1:
             query_dtype = (

vllm/model_executor/layers/attention/sparse_mla_attention.py

@@ -190,7 +190,8 @@ def __init__(
             pin_memory=PIN_MEMORY,
         )
         parallel_config = vllm_config.parallel_config
-        self.use_pcp = parallel_config.prefill_context_parallel_size > 1
+        self.pcp_world_size = parallel_config.prefill_context_parallel_size
+        self.use_pcp = self.pcp_world_size > 1
         try:
             self.dcp_world_size = get_dcp_group().world_size
         except AssertionError:
@@ -409,8 +410,12 @@ def build(
                 q_data_type=self.model_config.dtype,
                 output_dtype=self.model_config.dtype,
                 prefill_backend=self._prefill_backend,
-                use_dense_mha=_use_dense_mha_prefill(
-                    self.vllm_config, prefill_max_seq_len, self.topk_tokens
+                use_dense_mha=(
+                    _use_dense_mha_prefill(
+                        self.vllm_config, prefill_max_seq_len, self.topk_tokens
+                    )
+                    # Dense MHA cannot be used under PCP+DCP.
+                    and not (self.use_pcp and self.dcp_world_size > 1)
                 ),
                 topk_mask_workspace=self.topk_mask_workspace,
                 host_staging_plan=staging_plan,