74 lines
3.1 KiB
Python
74 lines
3.1 KiB
Python
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
# 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.
|
|
#
|
|
# Patch target: vllm.distributed.weight_transfer.factory.WeightTransferEngineFactory
|
|
#
|
|
# Replace the "nccl" and "ipc" factory entries with Ascend equivalents so that
|
|
# --weight-transfer-config '{"backend": "nccl"}' loads HCCLWeightTransferEngine
|
|
# and '{"backend": "ipc"}' loads NPUIPCWeightTransferEngine instead of the
|
|
# (unavailable) NCCL / CUDA IPC engines on Ascend NPU.
|
|
#
|
|
# Why this approach (factory swap) instead of patching Literal["nccl", "ipc"]:
|
|
# WeightTransferConfig.backend is a pydantic Literal["nccl", "ipc"].
|
|
# Adding "hccl" / "npu_ipc" would require modifying pydantic core schemas —
|
|
# fragile across pydantic versions. Swapping the factory entries means users
|
|
# pass the already-accepted "nccl" / "ipc" strings, but the factory resolves
|
|
# them to HCCL / NPU IPC.
|
|
#
|
|
# Timing — guaranteed to run before first factory usage:
|
|
#
|
|
# vllm serve main()
|
|
# line 24: from vllm.entrypoints.utils import ...
|
|
# → vllm.platforms.__getattr__("current_platform")
|
|
# → resolve_current_platform_cls_qualname()
|
|
# → vllm_ascend:register() → NPUPlatform()
|
|
# → NPUPlatform.pre_register_and_update()
|
|
# → adapt_patch(is_global_patch=True)
|
|
# → imports vllm_ascend.patch.platform
|
|
# → THIS PATCH RUNS ← "nccl" now points to HCCLWeightTransferEngine
|
|
# ...
|
|
# lines 82-86: subparser_init() → make_arg_parser()
|
|
# line 87: parse_args() → validates backend="nccl" via Literal (passes)
|
|
# ...
|
|
# later: worker init → WeightTransferEngineFactory.create_engine(config)
|
|
# → config.backend == "nccl" → factory loads HCCLWeightTransferEngine
|
|
#
|
|
# Future Plan:
|
|
# Remove this patch when upstream vllm relaxes the Literal type to str
|
|
# or provides an extension point for out-of-tree backends.
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
|
|
|
|
from vllm_ascend.distributed.weight_transfer.hccl_engine import (
|
|
HCCLWeightTransferEngine,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.distributed.weight_transfer.base import WeightTransferEngine
|
|
|
|
|
|
def _load_npu_ipc_engine() -> "type[WeightTransferEngine]":
|
|
from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import (
|
|
NPUIPCWeightTransferEngine,
|
|
)
|
|
|
|
return NPUIPCWeightTransferEngine
|
|
|
|
|
|
WeightTransferEngineFactory._registry["nccl"] = lambda: HCCLWeightTransferEngine
|
|
WeightTransferEngineFactory._registry["ipc"] = _load_npu_ipc_engine
|