b54d44e664
### What this PR does / why we need it? This PR adds the Prefill Context Parallelism (PCP) feature, which corresponds to DCP. For specific implementation details, please refer to the RFC https://github.com/vllm-project/vllm/issues/25749. TL;DR: PCP enhances long-sequence inference capabilities by partitioning the sequence dimension during the prefill stage. ### Does this PR introduce _any_ user-facing change? The current implementation primarily includes the following changes: Modified ModelRunner.py for CP partitioning logic for tokens; Modified attention_v1.py and mla_v1.py to adapt the GQA/MLA backend to PCP. Modified block_tables.py to extend the KV cache storage based on DCP&PCP; Added necessary command-line arguments to control parallelism for PCP; ### How was this patch tested? - vLLM version: v0.11.0rc3 - vLLM main: https://github.com/vllm-project/vllm/commit/v0.11.0 --------- Signed-off-by: LookAround <lixushi@huawei.com> Signed-off-by: chenjie <chenjie137@huawei.com> Signed-off-by: Delphine-Nic <tanwenqin@huawei.com> Signed-off-by: zhangsicheng5 <zhangsicheng5@huawei.com> Signed-off-by: Feng Liu <liufeng248@huawei.com> Signed-off-by: gaojc <1055866782@qq.com> Signed-off-by: weiguihua2 <weiguihua2@huawei.com> Signed-off-by: z50049692 <zhangmingwei11@huawei.com> Co-authored-by: chenjie <chenjie137@huawei.com> Co-authored-by: Delphine-Nic <tanwenqin@huawei.com> Co-authored-by: zhangsicheng5 <zhangsicheng5@huawei.com> Co-authored-by: Feng Liu <liufeng248@huawei.com> Co-authored-by: gaojc <1055866782@qq.com> Co-authored-by: weiguihua2 <weiguihua2@huawei.com> Co-authored-by: z50049692 <zhangmingwei11@huawei.com> Co-authored-by: w00896881 <wangzixuan40@huawei.com>
99 lines
3.7 KiB
Python
99 lines
3.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.
|
|
|
|
import os
|
|
import types
|
|
|
|
from tests.ut.kv_connector.utils import (create_request, create_scheduler,
|
|
create_vllm_config)
|
|
from vllm_ascend.distributed.llmdatadist_c_mgr_connector import (
|
|
LLMDataDistCMgrConnectorMetadata, LLMDataDistCMgrConnectorWorker, LLMRole)
|
|
|
|
|
|
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
|
|
|
|
|
|
def test_read_agent_metadata():
|
|
rank_table = {
|
|
"version":
|
|
"1.2",
|
|
"server_count":
|
|
"2",
|
|
"prefill_device_list": [{
|
|
"server_id": "192.168.1.1",
|
|
"device_id": "0",
|
|
"device_ip": "10.30.0.1",
|
|
"cluster_id": "0",
|
|
}, {
|
|
"server_id": "192.168.1.1",
|
|
"device_id": "1",
|
|
"device_ip": "10.30.0.2",
|
|
"cluster_id": "1",
|
|
}, {
|
|
"server_id": "192.168.1.2",
|
|
"device_id": "0",
|
|
"device_ip": "10.30.0.3",
|
|
"cluster_id": "2",
|
|
}, {
|
|
"server_id": "192.168.1.2",
|
|
"device_id": "1",
|
|
"device_ip": "10.30.0.4",
|
|
"cluster_id": "3",
|
|
}]
|
|
}
|
|
|
|
def get_device_ip(worker_local_ip, worker_tp_rank, worker_visible_devices):
|
|
old_visible_devices = os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "")
|
|
worker = types.SimpleNamespace()
|
|
worker.local_ip = worker_local_ip
|
|
worker.tp_rank = worker_tp_rank
|
|
worker.llm_datadist_role = LLMRole.PROMPT
|
|
worker.pcp_rank = 0
|
|
worker.tp_size = worker_tp_rank + 1
|
|
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = worker_visible_devices
|
|
agent_metadata = LLMDataDistCMgrConnectorWorker.read_agent_metadata(
|
|
worker, rank_table)
|
|
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = old_visible_devices
|
|
return agent_metadata.device_ip
|
|
|
|
assert get_device_ip("192.168.1.1", 0, "0") == "10.30.0.1"
|
|
assert get_device_ip("192.168.1.1", 0, "1") == "10.30.0.2"
|
|
assert get_device_ip("192.168.1.2", 0, "0") == "10.30.0.3"
|
|
assert get_device_ip("192.168.1.2", 0, "1") == "10.30.0.4"
|
|
assert get_device_ip("192.168.1.1", 0, "0,1") == "10.30.0.1"
|
|
assert get_device_ip("192.168.1.1", 1, "0,1") == "10.30.0.2"
|
|
assert get_device_ip("192.168.1.1", 0, "") == "10.30.0.1"
|
|
assert get_device_ip("192.168.1.1", 1, "") == "10.30.0.2"
|