237
tests/ut/_310p/test_block_table_310p.py
Normal file
237
tests/ut/_310p/test_block_table_310p.py
Normal file
@@ -0,0 +1,237 @@
|
||||
#
|
||||
# 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
|
||||
from vllm.distributed.parallel_state import GroupCoordinator
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
|
||||
|
||||
class TestBlockTable310(TestBase):
|
||||
def setUp(self):
|
||||
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]
|
||||
|
||||
def _create_block_table(self, dcp_world_size, dcp_rank, pcp_world_size, pcp_rank, cp_kv_cache_interleave_size):
|
||||
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 = 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 = 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._310p.block_table import BlockTable
|
||||
|
||||
return 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=0,
|
||||
)
|
||||
|
||||
def _create_multi_group_block_table(
|
||||
self,
|
||||
dcp_world_size,
|
||||
dcp_rank,
|
||||
pcp_world_size,
|
||||
pcp_rank,
|
||||
cp_kv_cache_interleave_size,
|
||||
block_sizes=None,
|
||||
max_num_blocks=None,
|
||||
kernel_sizes=None,
|
||||
):
|
||||
block_sizes = block_sizes or [self.block_size]
|
||||
max_num_blocks = max_num_blocks or [self.max_num_blocks_per_req] * len(block_sizes)
|
||||
kernel_sizes = kernel_sizes or [[self.block_size]] * len(block_sizes)
|
||||
|
||||
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 = 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 = 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._310p.block_table import MultiGroupBlockTable
|
||||
|
||||
return MultiGroupBlockTable(
|
||||
max_num_reqs=self.max_num_reqs,
|
||||
max_model_len=self.block_size * self.max_num_blocks_per_req,
|
||||
max_num_batched_tokens=self.max_num_batched_tokens,
|
||||
pin_memory=self.pin_memory,
|
||||
device=self.device,
|
||||
block_sizes=block_sizes,
|
||||
max_num_blocks=max_num_blocks,
|
||||
kernel_sizes=kernel_sizes,
|
||||
cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _setup_block_table_data(block_table, num_reqs=2):
|
||||
for i in range(num_reqs):
|
||||
block_ids = list(range(i * 4, (i + 1) * 4))
|
||||
block_table.add_row(block_ids, i)
|
||||
|
||||
def test_compute_slot_mapping_with_query_start_loc_signature(self):
|
||||
block_table = self._create_block_table(
|
||||
dcp_world_size=1,
|
||||
dcp_rank=0,
|
||||
pcp_world_size=1,
|
||||
pcp_rank=0,
|
||||
cp_kv_cache_interleave_size=1,
|
||||
)
|
||||
self._setup_block_table_data(block_table, num_reqs=2)
|
||||
|
||||
query_start_loc = torch.tensor([0, 2, 4], dtype=torch.int32)
|
||||
positions = torch.tensor([0, 1, 0, 1], dtype=torch.int64)
|
||||
|
||||
block_table.compute_slot_mapping(2, query_start_loc, positions)
|
||||
|
||||
expected = np.array([0, 1, 512, 513], dtype=np.int32)
|
||||
np.testing.assert_array_equal(block_table.slot_mapping.np[:4], expected)
|
||||
np.testing.assert_array_equal(block_table.slot_mapping.gpu[:4].cpu().numpy(), expected)
|
||||
|
||||
def test_multi_group_compute_slot_mapping_accepts_none_compressed_args(self):
|
||||
multi_group_block_table = self._create_multi_group_block_table(
|
||||
dcp_world_size=1,
|
||||
dcp_rank=0,
|
||||
pcp_world_size=1,
|
||||
pcp_rank=0,
|
||||
cp_kv_cache_interleave_size=1,
|
||||
)
|
||||
self._setup_block_table_data(multi_group_block_table[0], num_reqs=2)
|
||||
|
||||
query_start_loc = torch.tensor([0, 2, 4], dtype=torch.int32)
|
||||
positions = torch.tensor([0, 1, 0, 1], dtype=torch.int64)
|
||||
|
||||
multi_group_block_table.compute_slot_mapping(2, query_start_loc, positions, None, None)
|
||||
|
||||
expected = np.array([0, 1, 512, 513], dtype=np.int32)
|
||||
np.testing.assert_array_equal(multi_group_block_table[0].slot_mapping.np[:4], expected)
|
||||
np.testing.assert_array_equal(multi_group_block_table[0].slot_mapping.gpu[:4].cpu().numpy(), expected)
|
||||
|
||||
def test_multi_group_compute_slot_mapping_uses_compressed_inputs_per_group(self):
|
||||
multi_group_block_table = self._create_multi_group_block_table(
|
||||
dcp_world_size=1,
|
||||
dcp_rank=0,
|
||||
pcp_world_size=1,
|
||||
pcp_rank=0,
|
||||
cp_kv_cache_interleave_size=1,
|
||||
block_sizes=[self.block_size, self.block_size],
|
||||
max_num_blocks=[self.max_num_blocks_per_req, self.max_num_blocks_per_req],
|
||||
kernel_sizes=[[self.block_size], [self.block_size]],
|
||||
)
|
||||
for block_table in multi_group_block_table.block_tables:
|
||||
self._setup_block_table_data(block_table, num_reqs=2)
|
||||
|
||||
query_start_loc = torch.tensor([0, 2, 4], dtype=torch.int32)
|
||||
positions = torch.tensor([0, 1, 0, 1], dtype=torch.int64)
|
||||
positions_compressed_list = [
|
||||
np.array([0, 1], dtype=np.int64),
|
||||
np.array([0], dtype=np.int64),
|
||||
]
|
||||
req_indices_compressed_list = [
|
||||
np.array([0, 0], dtype=np.int64),
|
||||
np.array([1], dtype=np.int64),
|
||||
]
|
||||
|
||||
multi_group_block_table.compute_slot_mapping(
|
||||
2,
|
||||
query_start_loc,
|
||||
positions,
|
||||
positions_compressed_list,
|
||||
req_indices_compressed_list,
|
||||
)
|
||||
|
||||
np.testing.assert_array_equal(
|
||||
multi_group_block_table[0].slot_mapping.np[:2],
|
||||
np.array([0, 1], dtype=np.int32),
|
||||
)
|
||||
np.testing.assert_array_equal(
|
||||
multi_group_block_table[1].slot_mapping.np[:1],
|
||||
np.array([512], dtype=np.int32),
|
||||
)
|
||||
|
||||
def test_compute_slot_mapping_with_req_indices_signature(self):
|
||||
block_table = self._create_block_table(
|
||||
dcp_world_size=4,
|
||||
dcp_rank=0,
|
||||
pcp_world_size=2,
|
||||
pcp_rank=0,
|
||||
cp_kv_cache_interleave_size=1,
|
||||
)
|
||||
self._setup_block_table_data(block_table, num_reqs=1)
|
||||
|
||||
req_indices = np.zeros(16, dtype=np.int32)
|
||||
positions = np.arange(16, dtype=np.int32)
|
||||
|
||||
block_table.compute_slot_mapping(req_indices, positions)
|
||||
|
||||
expected = np.array([0, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, -1, -1, -1, -1], dtype=np.int32)
|
||||
np.testing.assert_array_equal(block_table.slot_mapping.np[:16], expected)
|
||||
np.testing.assert_array_equal(block_table.slot_mapping.gpu[:16].cpu().numpy(), expected)
|
||||
|
||||
def test_compute_slot_mapping_rejects_device_tensor_inputs(self):
|
||||
block_table = self._create_block_table(
|
||||
dcp_world_size=1,
|
||||
dcp_rank=0,
|
||||
pcp_world_size=1,
|
||||
pcp_rank=0,
|
||||
cp_kv_cache_interleave_size=1,
|
||||
)
|
||||
self._setup_block_table_data(block_table, num_reqs=2)
|
||||
|
||||
req_indices = np.array([0, 0, 1, 1], dtype=np.int64)
|
||||
device_positions = torch.empty(4, dtype=torch.int64, device="meta")
|
||||
|
||||
with self.assertRaisesRegex(TypeError, "D2H"):
|
||||
block_table.compute_slot_mapping(req_indices, device_positions)
|
||||
|
||||
device_query_start_loc = torch.empty(3, dtype=torch.int32, device="meta")
|
||||
positions = torch.arange(4, dtype=torch.int64)
|
||||
|
||||
with self.assertRaisesRegex(TypeError, "D2H"):
|
||||
block_table.compute_slot_mapping(2, device_query_start_loc, positions)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user