[Misc] Fix import error and address nits to make CI happy (#563)

1. Add `vllm_version_is` function to check vllm version.
2. `ensure_kv_transfer_initialized` and `get_kv_transfer_group ` have
been moved to other place in vllm main branch via
3408e47159
, this patch fix the import error.

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
This commit is contained in:
wangxiyuan
2025-04-18 12:23:32 +08:00
committed by GitHub
parent 66a0837963
commit 42c7fbb10e
9 changed files with 46 additions and 30 deletions

View File

@@ -25,8 +25,7 @@ import torch.distributed
from torch import nn
from vllm import envs
from vllm.config import VllmConfig
from vllm.distributed import (ensure_kv_transfer_initialized,
ensure_model_parallel_initialized,
from vllm.distributed import (ensure_model_parallel_initialized,
init_distributed_environment,
set_custom_all_reduce)
from vllm.logger import logger
@@ -45,10 +44,15 @@ from vllm.worker.worker_base import (LocalOrDistributedWorkerBase, WorkerBase,
WorkerInput)
from vllm_ascend.platform import NPUPlatform
from vllm_ascend.utils import try_register_lib
from vllm_ascend.utils import try_register_lib, vllm_version_is
from vllm_ascend.worker.model_runner import NPUModelRunner
from vllm_ascend.worker.pooling_model_runner import NPUPoolingModelRunner
if vllm_version_is("0.8.4"):
from vllm.distributed import ensure_kv_transfer_initialized
else:
from vllm.distributed.kv_transfer import ensure_kv_transfer_initialized
class NPUWorker(LocalOrDistributedWorkerBase):
"""A worker class that executes (a partition of) the model on a NPU.