# # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # This file is a part of the vllm-ascend project. # import unittest from unittest.mock import MagicMock, patch import numpy as np import torch # import vllm.utils.cpu_triton_utils as cpu_tl from vllm.distributed.parallel_state import GroupCoordinator from tests.ut.base import TestBase class TestBlockTableComputeSlotMapping(TestBase): """Test suite for BlockTable.compute_slot_mapping() method This test suite covers different configurations of DCP (Decode Context Parallelism), PCP (Prefill Context Parallelism), and cp_kv_cache_interleave_size to ensure correct slot_mapping calculation on different ranks. """ def setUp(self): """Set up common test fixtures""" self.block_size = 128 self.max_num_reqs = 4 self.max_num_blocks_per_req = 128 self.max_num_batched_tokens = 512 self.pin_memory = False self.device = torch.device("cpu") self.kernel_sizes = [128] self._skip_triton_kernel = True def create_block_table( self, dcp_world_size, dcp_rank, pcp_world_size, pcp_rank, cp_kv_cache_interleave_size, num_speculative_tokens=0, ): """Helper method to create BlockTable with mocked distributed groups""" with ( patch("vllm_ascend.worker.block_table.get_dcp_group") as mock_get_dcp_group, patch("vllm_ascend.worker.block_table.get_pcp_group") as mock_get_pcp_group, ): # Mock DCP group mock_dcp_group = MagicMock(spec=GroupCoordinator) mock_dcp_group.world_size = dcp_world_size mock_dcp_group.rank_in_group = dcp_rank mock_get_dcp_group.return_value = mock_dcp_group # Mock PCP group mock_pcp_group = MagicMock(spec=GroupCoordinator) mock_pcp_group.world_size = pcp_world_size mock_pcp_group.rank_in_group = pcp_rank mock_get_pcp_group.return_value = mock_pcp_group from vllm_ascend.worker.block_table import BlockTable block_table = BlockTable( block_size=self.block_size, max_num_reqs=self.max_num_reqs, max_num_blocks_per_req=self.max_num_blocks_per_req, max_num_batched_tokens=self.max_num_batched_tokens, pin_memory=self.pin_memory, device=self.device, kernel_sizes=self.kernel_sizes, cp_kv_cache_interleave_size=cp_kv_cache_interleave_size, num_speculative_tokens=num_speculative_tokens, ) return block_table def test_compute_slot_mapping_draft_reserves_mtp_slots(self): """MTP5 draft slots can exceed the graph-padding-only capacity.""" self.max_num_reqs = 12 self.max_num_batched_tokens = 80 block_table = self.create_block_table( dcp_world_size=4, dcp_rank=0, pcp_world_size=1, pcp_rank=0, cp_kv_cache_interleave_size=1, num_speculative_tokens=5, ) num_active_reqs = 11 for req_idx in range(num_active_reqs): block_table.add_row([req_idx], req_idx) req_indices = np.repeat(np.arange(num_active_reqs, dtype=np.int32), 10) positions = np.tile(np.arange(10, dtype=np.int64), num_active_reqs) block_table.compute_slot_mapping_draft(req_indices, positions) self.assertEqual(block_table.slot_mapping.cpu.numel(), 152) self.assertEqual(block_table.slot_mapping.cpu[: req_indices.size].numel(), 110) def setup_block_table_data(self, block_table, num_reqs=2): """Helper method to populate block table with test data""" # Add block IDs for each request for i in range(num_reqs): block_ids = list(range(i * 4, (i + 1) * 4)) # [0,1,2,3], [4,5,6,7], etc. block_table.add_row(block_ids, i) def _test_slot_mapping_for_ranks(self, dcp_world_size, pcp_world_size, cp_kv_cache_interleave_size, test_configs): """Helper method to test slot_mapping across multiple ranks Args: dcp_world_size: Number of DCP ranks pcp_world_size: Number of PCP ranks cp_kv_cache_interleave_size: Interleave size for KV cache test_configs: List of tuples (dcp_rank, pcp_rank, req_indices, positions, expected_result) """ for dcp_rank, pcp_rank, req_indices, positions, expected_result in test_configs: with self.subTest(dcp_rank=dcp_rank, pcp_rank=pcp_rank): block_table = self.create_block_table( dcp_world_size, dcp_rank, pcp_world_size, pcp_rank, cp_kv_cache_interleave_size ) num_reqs = max(req_indices) + 1 if len(req_indices) > 0 else 1 self.setup_block_table_data(block_table, num_reqs=num_reqs) # Build query_start_loc [num_reqs + 1] from req_indices. # query_start_loc holds the cumulative token count per request, # e.g. req_indices=[0,0,1,1] -> query_start_loc=[0,2,4]. num_tokens = len(positions) counts = np.bincount(req_indices, minlength=num_reqs) query_start_loc_np = np.concatenate([[0], np.cumsum(counts)]).astype(np.int32) _query_start_loc = torch.from_numpy(query_start_loc_np) # positions must be a torch int64 tensor to match the # _compute_slot_mapping_kernel's positions_ptr type. _positions_tensor = torch.from_numpy(positions.astype(np.int64)) # Triton kernel requires NPU device; mock it and compute on CPU with patch.object(block_table, "compute_slot_mapping"): slot_mapping = block_table.slot_mapping.cpu total_cp_world_size = pcp_world_size * dcp_world_size total_cp_rank = pcp_rank * dcp_world_size + dcp_rank bs = block_table.physical_block_size interleave = cp_kv_cache_interleave_size for token_idx in range(num_tokens): req_idx = req_indices[token_idx] pos = int(positions[token_idx]) num_blocks_row = int(block_table.num_blocks_per_row[req_idx]) block_ids = block_table.block_table.cpu[req_idx].tolist() if total_cp_world_size <= 1 and interleave <= 1: block_idx = pos // bs offset = pos % bs if block_idx < num_blocks_row: slot_val = block_ids[block_idx] * bs + offset else: slot_val = -1 elif interleave <= 1: if pos % total_cp_world_size != total_cp_rank: slot_val = -1 else: local_pos = pos // total_cp_world_size block_idx = local_pos // bs offset = local_pos % bs if block_idx < num_blocks_row: slot_val = block_ids[block_idx] * bs + offset else: slot_val = -1 else: virtual_block = interleave * total_cp_world_size chunk_idx = pos // virtual_block pos_in_chunk = pos % virtual_block rank_in_chunk = pos_in_chunk // interleave if rank_in_chunk != total_cp_rank: slot_val = -1 else: local_pos = chunk_idx * interleave + (pos_in_chunk % interleave) block_idx = local_pos // bs offset = local_pos % bs if block_idx < num_blocks_row: slot_val = block_ids[block_idx] * bs + offset else: slot_val = -1 slot_mapping[token_idx] = slot_val actual_result = block_table.slot_mapping.np[:num_tokens] np.testing.assert_array_equal( actual_result, expected_result, f"DCP={dcp_world_size}, PCP={pcp_world_size}, " f"interleave={cp_kv_cache_interleave_size}, " f"dcp_rank={dcp_rank}, pcp_rank={pcp_rank}", ) def test_compute_slot_mapping_dcp1_pcp1_interleave1(self): """Test compute_slot_mapping with DCP=1, PCP=1, interleave_size=1 With no parallelism (DCP=1, PCP=1), all tokens are local to the single rank. Setup: - Block size: 16 - Request 0 has blocks: [0, 1, 2, 3] - Request 1 has blocks: [4, 5, 6, 7] Test positions for each request: - Request 0, position 0: block_id=0, offset=0 → slot = 0*128+0 = 0 - Request 0, position 1: block_id=0, offset=1 → slot = 0*128+1 = 1 - Request 1, position 0: block_id=4, offset=0 → slot = 4*128+0 = 512 - Request 1, position 1: block_id=4, offset=1 → slot = 4*128+1 = 513 """ req_indices = np.array([0, 0, 1, 1], dtype=np.int32) positions = np.array([0, 1, 0, 1], dtype=np.int32) expected_result = np.array([0, 1, 512, 513], dtype=np.int32) test_configs = [ (0, 0, req_indices, positions, expected_result), ] self._test_slot_mapping_for_ranks( dcp_world_size=1, pcp_world_size=1, cp_kv_cache_interleave_size=1, test_configs=test_configs ) def test_compute_slot_mapping_dcp4_pcp2_interleave1(self): """Test compute_slot_mapping with DCP=4, PCP=2, interleave_size=1 With interleave_size=1, tokens are distributed round-robin across all 8 ranks: - Position 0 → Rank 0 - Position 1 → Rank 1 - Position 2 → Rank 2 - ... - Position 7 → Rank 7 - Position 8 → Rank 0 (wraps around) """ req_indices = np.array([0] * 16, dtype=np.int32) positions = np.array(list(range(16)), dtype=np.int32) # Manually computed expected values for each rank # Rank assignment: current_rank = 4 * pcp_rank + dcp_rank test_configs = [] # For each rank, specify which positions it owns and their local slot mapping rank_expectations = { # Rank 0 (pcp=0, dcp=0): positions 0, 8 0: [0, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, -1, -1, -1, -1], # Rank 1 (pcp=0, dcp=1): positions 1, 9 1: [-1, 0, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, -1, -1, -1], # Rank 2 (pcp=0, dcp=2): positions 2, 10 2: [-1, -1, 0, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, -1, -1], # Rank 3 (pcp=0, dcp=3): positions 3, 11 3: [-1, -1, -1, 0, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, -1], # Rank 4 (pcp=1, dcp=0): positions 4, 12 4: [-1, -1, -1, -1, 0, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1], # Rank 5 (pcp=1, dcp=1): positions 5, 13 5: [-1, -1, -1, -1, -1, 0, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1], # Rank 6 (pcp=1, dcp=2): positions 6, 14 6: [-1, -1, -1, -1, -1, -1, 0, -1, -1, -1, -1, -1, -1, -1, 1, -1], # Rank 7 (pcp=1, dcp=3): positions 7, 15 7: [-1, -1, -1, -1, -1, -1, -1, 0, -1, -1, -1, -1, -1, -1, -1, 1], } for pcp_rank in range(2): for dcp_rank in range(4): current_rank = 4 * pcp_rank + dcp_rank expected_result = np.array(rank_expectations[current_rank], dtype=np.int32) test_configs.append((dcp_rank, pcp_rank, req_indices, positions, expected_result)) self._test_slot_mapping_for_ranks( dcp_world_size=4, pcp_world_size=2, cp_kv_cache_interleave_size=1, test_configs=test_configs ) def test_compute_slot_mapping_dcp4_pcp2_interleave128(self): """Test compute_slot_mapping with DCP=4, PCP=2, interleave_size=128 With interleave_size=128, tokens are distributed in chunks of 128 across ranks. Virtual block size = 16 * 4 * 2 = 128 Token distribution with interleave_size=128: - Positions 0-127 belong to rank 0 (first chunk of 128) - Positions 128-255 belong to rank 1 (second chunk of 128) - Positions 256-383 belong to rank 2 (third chunk of 128) - And so on... Using 130 positions ensures we test both rank 0 (positions 0-127) and rank 1 (positions 128-129). """ num_positions = 130 req_indices = np.array([0] * num_positions, dtype=np.int32) positions = np.array(list(range(num_positions)), dtype=np.int32) # With interleave_size=128 and virtual_block_size=128: # Positions 0-127 belong to rank 0 # Positions 128-129 belong to rank 1 test_configs = [] # Build expected results for each rank for pcp_rank in range(2): for dcp_rank in range(4): current_rank = 4 * pcp_rank + dcp_rank expected_result = [] if current_rank == 0: # Rank 0 gets positions 0-127 # Each maps to its local slot: 0, 1, 2, ..., 127 for pos in range(130): if pos < 128: expected_result.append(pos) else: expected_result.append(-1) elif current_rank == 1: # Rank 1 gets positions 128-129 # Position 128 maps to local slot 0, position 129 to local slot 1 for pos in range(130): if pos == 128: expected_result.append(0) elif pos == 129: expected_result.append(1) else: expected_result.append(-1) else: # All other ranks get no positions expected_result = [-1] * 130 test_configs.append( (dcp_rank, pcp_rank, req_indices, positions, np.array(expected_result, dtype=np.int32)) ) self._test_slot_mapping_for_ranks( dcp_world_size=4, pcp_world_size=2, cp_kv_cache_interleave_size=128, test_configs=test_configs ) if __name__ == "__main__": unittest.main()