[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>
Lucas Wilkinson committed
Sep 13, 2026 at 16:45 UTC
586f652f8d9e841901c40634c024a2ebbfe3402f
2 files changed
+40
-13
tests/v1/attention/test_indexer_dcp_localize.py
+32
-11
@@ -966,34 +966,55 @@ def test_sparse_decode_dcp_short_context_matches_non_dcp():
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"""
969
+@pytest.mark.parametrize("interleave", [1, 64])
970
+@pytest.mark.parametrize("rows_per_req", [1, 2])
971
+@pytest.mark.parametrize(
972
+ "req_lens", [[7], [8, 8], [1, 9], [63, 64, 65], [255, 256, 257], [256, 1, 2730]]
973
+)
974
+def test_pcp_plan_deinterleave_restores_global_order(
975
+ monkeypatch, dcp_world_size, interleave, rows_per_req, req_lens
976
+):
977
+ """Gathering block-interleaved shards must restore each request's token order."""
978
from vllm.v1.attention.backends.mla.indexer import build_pcp_global_chunk_plan
979
980
+ monkeypatch.setattr("vllm.v1.attention.backends.mla.indexer.PIN_MEMORY", False)
981
+ monkeypatch.setattr(
982
+ "vllm.v1.attention.backends.mla.indexer.async_copy_to_gpu",
983
+ lambda x, device: x.to(device),
984
+ )
985
scheduled = np.array(req_lens, dtype=np.int64)
975
- shard_rows = -(-scheduled // dcp_world_size)
976
- rows = np.arange(len(scheduled))
986
+ shard_rows = np.array(
987
+ [
988
+ max(
989
+ _local_count(g, r, dcp_world_size, interleave)
990
+ for r in range(dcp_world_size)
991
+ )
992
+ for g in req_lens
993
+ ]
994
+ )
995
+ rows = np.repeat(np.arange(len(scheduled)), rows_per_req)
996
plan = build_pcp_global_chunk_plan(
978
- rows, shard_rows, dcp_world_size, torch.device("cpu")
997
+ rows, shard_rows[rows], dcp_world_size, torch.device("cpu"), interleave
998
)
999
1000
starts = np.concatenate([[0], np.cumsum(scheduled)])
1001
1002
# Build each rank's padded shard exactly as the cache gather would.
984
- padded_cu = plan.padded_local_cu.tolist()
1003
+ padded_cu = plan.padded_local_cu[::rows_per_req].tolist()
1004
shards = torch.zeros(dcp_world_size, plan.padded_local_total, 1)
1005
for r in range(dcp_world_size):
1006
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
1007
+ owned = [t for t in range(g) if (t // interleave) % dcp_world_size == r]
1008
+ for local, t in enumerate(owned):
1009
+ shards[r, padded_cu[i] + local, 0] = starts[i] + t
1010
1011
gathered = shards.reshape(dcp_world_size * plan.padded_local_total, 1)
1012
restored = gathered[plan.deinterleave_idx]
1013
# The layout is padded per request; only the real positions are read.
1014
row_start = plan.row_start_cu.tolist()
995
- for i, g in enumerate(req_lens):
1015
+ for row, i in enumerate(rows):
1016
+ g = req_lens[i]
1017
torch.testing.assert_close(
997
- restored[row_start[i] : row_start[i] + g, 0],
1018
+ restored[row_start[row] : row_start[row] + g, 0],
1019
torch.arange(starts[i], starts[i] + g, dtype=torch.float32),
1020
)
vllm/v1/attention/backends/mla/indexer.py
+8
-2
@@ -25,6 +25,7 @@ from vllm.utils.deep_gemm import (
25
has_deep_gemm,
26
native_next_n_supported,
27
)
28
+from vllm.utils.math_utils import round_down
29
from vllm.utils.platform_utils import num_compute_units
30
from vllm.utils.torch_utils import PIN_MEMORY
31
from vllm.v1.attention.backend import (
@@ -286,6 +287,7 @@ def build_pcp_global_chunk_plan(
287
row_shard_rows: np.ndarray,
288
dcp_world_size: int,
289
device: torch.device,
290
+ interleave: int = 1,
291
) -> PCPGlobalChunkPlan:
292
"""Plan the PCP packing for one chunk from its requests' KV shard rows.
293
@@ -327,10 +329,13 @@ def build_pcp_global_chunk_plan(
329
for i in range(len(region_first_row)):
330
g = int(region_padded[i]) * dcp_world_size
331
t = np.arange(g, dtype=np.int64)
332
+ local = round_down(t // dcp_world_size, interleave) + t % interleave
333
+ # Padded positions are outside causal bounds but must remain in-bounds.
334
+ local = np.minimum(local, region_padded[i] - 1)
335
idx_np[region_start[i] : region_start[i] + g] = (
331
- (t % dcp_world_size) * padded_total
336
+ ((t // interleave) % dcp_world_size) * padded_total
337
+ region_padded_start[i]
333
- + t // dcp_world_size
338
+ + local
339
)
340
341
cu = async_copy_to_gpu(cu, device=device)
@@ -1268,6 +1273,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1273
shard_rows[req_slice],
1274
self.dcp_world_size,
1275
self.device,
1276
+ self.cp_kv_cache_interleave_size,
1277
)
1278
metadata = build_prefill_chunk_metadata(
1279
req_slice.start,