43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
|
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||
|
|
|
||
|
|
from tests.ut.kv_connector.utils import (create_request, create_scheduler,
|
||
|
|
create_vllm_config)
|
||
|
|
from vllm_ascend.distributed.llmdatadist_c_mgr_connector import \
|
||
|
|
LLMDataDistCMgrConnectorMetadata
|
||
|
|
|
||
|
|
|
||
|
|
def test_basic_inferface():
|
||
|
|
"""Unit test for basic LLMDataDistCMgrConnector interface functionality."""
|
||
|
|
|
||
|
|
vllm_config = create_vllm_config()
|
||
|
|
scheduler = create_scheduler(vllm_config)
|
||
|
|
|
||
|
|
# 2 Full Blocks and 1 Half Block.
|
||
|
|
BLOCK_SIZE = vllm_config.cache_config.block_size
|
||
|
|
NUM_EXTERNAL_FULL_BLOCKS = 2
|
||
|
|
NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
|
||
|
|
|
||
|
|
request = create_request(request_id=1,
|
||
|
|
num_tokens=NUM_TOKENS,
|
||
|
|
do_remote_prefill=True)
|
||
|
|
request_id = request.request_id
|
||
|
|
|
||
|
|
scheduler.add_request(request)
|
||
|
|
|
||
|
|
# Remote Prefill, triggers LLMDataDistCMgrConnectorMetadata.
|
||
|
|
scheduler_output = scheduler.schedule()
|
||
|
|
kv_connector_metadata = scheduler_output.kv_connector_metadata
|
||
|
|
assert kv_connector_metadata is not None
|
||
|
|
assert isinstance(kv_connector_metadata, LLMDataDistCMgrConnectorMetadata)
|
||
|
|
|
||
|
|
assert len(kv_connector_metadata.requests) == 1
|
||
|
|
assert request_id in kv_connector_metadata.requests
|
||
|
|
req_meta = kv_connector_metadata.requests[request_id]
|
||
|
|
|
||
|
|
for block_id, block in zip(
|
||
|
|
req_meta.local_block_ids, scheduler.kv_cache_manager.coordinator.
|
||
|
|
single_type_managers[0].req_to_blocks[request_id]):
|
||
|
|
assert block_id == block.block_id
|