[Bugfix][PCP][DCP] Respect interleave in indexer KV gather mapping (#56715)

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

tests/v1/attention/test_indexer_dcp_localize.py

@@ -966,34 +966,55 @@ def test_sparse_decode_dcp_short_context_matches_non_dcp():
 
 
 @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"""
+@pytest.mark.parametrize("interleave", [1, 64])
+@pytest.mark.parametrize("rows_per_req", [1, 2])
+@pytest.mark.parametrize(
+    "req_lens", [[7], [8, 8], [1, 9], [63, 64, 65], [255, 256, 257], [256, 1, 2730]]
+)
+def test_pcp_plan_deinterleave_restores_global_order(
+    monkeypatch, dcp_world_size, interleave, rows_per_req, req_lens
+):
+    """Gathering block-interleaved shards must restore each request's token order."""
     from vllm.v1.attention.backends.mla.indexer import build_pcp_global_chunk_plan
 
+    monkeypatch.setattr("vllm.v1.attention.backends.mla.indexer.PIN_MEMORY", False)
+    monkeypatch.setattr(
+        "vllm.v1.attention.backends.mla.indexer.async_copy_to_gpu",
+        lambda x, device: x.to(device),
+    )
     scheduled = np.array(req_lens, dtype=np.int64)
-    shard_rows = -(-scheduled // dcp_world_size)
-    rows = np.arange(len(scheduled))
+    shard_rows = np.array(
+        [
+            max(
+                _local_count(g, r, dcp_world_size, interleave)
+                for r in range(dcp_world_size)
+            )
+            for g in req_lens
+        ]
+    )
+    rows = np.repeat(np.arange(len(scheduled)), rows_per_req)
     plan = build_pcp_global_chunk_plan(
-        rows, shard_rows, dcp_world_size, torch.device("cpu")
+        rows, shard_rows[rows], dcp_world_size, torch.device("cpu"), interleave
     )
 
     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()
+    padded_cu = plan.padded_local_cu[::rows_per_req].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
+            owned = [t for t in range(g) if (t // interleave) % dcp_world_size == r]
+            for local, t in enumerate(owned):
+                shards[r, padded_cu[i] + local, 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):
+    for row, i in enumerate(rows):
+        g = req_lens[i]
         torch.testing.assert_close(
-            restored[row_start[i] : row_start[i] + g, 0],
+            restored[row_start[row] : row_start[row] + g, 0],
             torch.arange(starts[i], starts[i] + g, dtype=torch.float32),
         )

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

@@ -25,6 +25,7 @@ from vllm.utils.deep_gemm import (
     has_deep_gemm,
     native_next_n_supported,
 )
+from vllm.utils.math_utils import round_down
 from vllm.utils.platform_utils import num_compute_units
 from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.attention.backend import (
@@ -286,6 +287,7 @@ def build_pcp_global_chunk_plan(
     row_shard_rows: np.ndarray,
     dcp_world_size: int,
     device: torch.device,
+    interleave: int = 1,
 ) -> PCPGlobalChunkPlan:
     """Plan the PCP packing for one chunk from its requests' KV shard rows.
 
@@ -327,10 +329,13 @@ def build_pcp_global_chunk_plan(
     for i in range(len(region_first_row)):
         g = int(region_padded[i]) * dcp_world_size
         t = np.arange(g, dtype=np.int64)
+        local = round_down(t // dcp_world_size, interleave) + t % interleave
+        # Padded positions are outside causal bounds but must remain in-bounds.
+        local = np.minimum(local, region_padded[i] - 1)
         idx_np[region_start[i] : region_start[i] + g] = (
-            (t % dcp_world_size) * padded_total
+            ((t // interleave) % dcp_world_size) * padded_total
             + region_padded_start[i]
-            + t // dcp_world_size
+            + local
         )
 
     cu = async_copy_to_gpu(cu, device=device)
@@ -1268,6 +1273,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
                         shard_rows[req_slice],
                         self.dcp_world_size,
                         self.device,
+                        self.cp_kv_cache_interleave_size,
                     )
                 metadata = build_prefill_chunk_metadata(
                     req_slice.start,