From 586f652f8d9e841901c40634c024a2ebbfe3402f Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sun, 13 Sep 2026 16:45:17 -0400 Subject: [PATCH] [Bugfix][PCP][DCP] Respect interleave in indexer KV gather mapping (#56715) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- .../v1/attention/test_indexer_dcp_localize.py | 43 ++++++++++++++----- vllm/v1/attention/backends/mla/indexer.py | 10 ++++- 2 files changed, 40 insertions(+), 13 deletions(-) diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index ffcfcbfac1..eab4759640 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/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), ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 164a3f8916..828f2c2ac4 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/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,