From e6004e30943b443048a996867f964fb1dad23951 Mon Sep 17 00:00:00 2001 From: recky-c <252898884+recky-c@users.noreply.github.com> Date: Thu, 17 Sep 2026 14:57:09 +0800 Subject: [PATCH 1/2] fix: match PCP slot buffer dtype to Ascend block tables Align persistent global and gathered PCP slot buffers and the padding scalar with the normal block-table dtype during initialization. Add gather and dummy buffer regression coverage. Validation: Python syntax and git diff --check passed. Runtime tests and format.sh ci remain unverified because PyTorch, pytest and pre-commit are unavailable. AI assistance: Codex. Signed-off-by: recky-c <252898884+recky-c@users.noreply.github.com> --- tests/ut/worker/test_pcp_manager_v2.py | 30 ++++++++++++++++++++++++++ vllm_ascend/worker/v2/pcp_manager.py | 13 +++++++++++ 2 files changed, 43 insertions(+) diff --git a/tests/ut/worker/test_pcp_manager_v2.py b/tests/ut/worker/test_pcp_manager_v2.py index 7c0e3923dac9..5b28e9618364 100644 --- a/tests/ut/worker/test_pcp_manager_v2.py +++ b/tests/ut/worker/test_pcp_manager_v2.py @@ -1039,3 +1039,33 @@ def test_dummy_attention_context_uses_current_batch(pcp_rank, has_stale_batch): assert manager._global_batch is saved_batch assert manager._hidden_restore_idx is saved_indices manager._block_tables.gather_block_tables.assert_not_called() + + +@pytest.mark.parametrize("slot_dtype", [torch.int32, torch.int64]) +def test_pcp_slot_buffers_match_block_tables(slot_dtype): + def init_buffers(manager, **kwargs): + manager._global_batch_slot_mappings = torch.empty((1, 4), dtype=torch.int64) + manager._gathered_kv_slot_mappings = torch.empty((1, 8), dtype=torch.int64) + manager._pad_slot_id = torch.tensor(-1, dtype=torch.int64) + manager.pcp_world_size = 2 + + block_tables = SimpleNamespace(slot_mappings=torch.empty((1, 4), dtype=slot_dtype)) + with patch.object(PCPManager, "__init__", init_buffers): + manager = AscendPCPManager(2, 0, torch.device("cpu"), block_tables=block_tables) + + assert manager._global_batch_slot_mappings.dtype == slot_dtype + assert manager._gathered_kv_slot_mappings.dtype == slot_dtype + manager._global_batch_slot_mappings.copy_(torch.tensor([[3200, 3201, -1, 3203]])) + manager._padded_gather_idx = torch.tensor([3, 0, 1, 2]) + manager._gathered_kv_write_mask = torch.tensor([True, True, False, True]) + buffer_ptr = manager._gathered_kv_slot_mappings.data_ptr() + + gathered = manager._convert_to_gathered_slot_mappings(manager._global_batch_slot_mappings) + assert gathered.tolist() == [[3203, 3200, -1, -1]] + assert gathered.dtype == slot_dtype + assert gathered.data_ptr() == buffer_ptr + + dummy = manager.get_dummy_slot_mappings(2) + assert dummy.tolist() == [[-1, -1, -1, -1]] + assert dummy.dtype == slot_dtype + assert dummy.data_ptr() == buffer_ptr \ No newline at end of file diff --git a/vllm_ascend/worker/v2/pcp_manager.py b/vllm_ascend/worker/v2/pcp_manager.py index daf250c73d64..030eb451cd5d 100644 --- a/vllm_ascend/worker/v2/pcp_manager.py +++ b/vllm_ascend/worker/v2/pcp_manager.py @@ -77,6 +77,19 @@ def __init__( cp_interleave=cp_interleave, ) + # Match the normal Ascend block-table path before any graph capture. + if block_tables is not None: + slot_dtype = block_tables.slot_mappings.dtype + if self._global_batch_slot_mappings is not None: + self._global_batch_slot_mappings = torch.empty_like( + self._global_batch_slot_mappings, dtype=slot_dtype + ) + if self._gathered_kv_slot_mappings is not None: + self._gathered_kv_slot_mappings = torch.empty_like( + self._gathered_kv_slot_mappings, dtype=slot_dtype + ) + self._pad_slot_id = self._pad_slot_id.to(slot_dtype) + # vLLM #53515 made the PCP-local buffers persistent and uses them for # graph capture. Preserve that ownership while providing the extra CPU # and NumPy sequence-length views required by AscendInputBatch. From 0ecafb173857af82c9c7cd830123d8f9f979d4c9 Mon Sep 17 00:00:00 2001 From: recky-c <252898884+recky-c@users.noreply.github.com> Date: Thu, 17 Sep 2026 14:59:32 +0800 Subject: [PATCH 2/2] docs: clarify why PCP slot buffers follow block-table dtype AI assistance: Codex. Comment-only change; git diff --check passed. format.sh ci is blocked by missing pre-commit. Signed-off-by: recky-c <252898884+recky-c@users.noreply.github.com> --- vllm_ascend/worker/v2/pcp_manager.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm_ascend/worker/v2/pcp_manager.py b/vllm_ascend/worker/v2/pcp_manager.py index 030eb451cd5d..7f173c4d1093 100644 --- a/vllm_ascend/worker/v2/pcp_manager.py +++ b/vllm_ascend/worker/v2/pcp_manager.py @@ -77,7 +77,8 @@ def __init__( cp_interleave=cp_interleave, ) - # Match the normal Ascend block-table path before any graph capture. + # PCP supplies its own output buffers to compute_slot_mappings, so their + # dtype must match Ascend block-table slots for cache-write operators. if block_tables is not None: slot_dtype = block_tables.slot_mappings.dtype if self._global_batch_slot_mappings is not None: