EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[Worker]Lazy import torch_npu (#184)

Browse Source 
### What this PR does / why we need it?
To avoid unnecessary delays, we only import torch_npu when profilling is
enabled.

Signed-off-by: wangli <wangli858794774@gmail.com>
This commit is contained in:
Li Wang
2025-02-27 16:52:11 +08:00
committed by GitHub
parent 6042c210bc
commit c131e43e7d
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
vllm_ascend/worker.py
View File

@@ -22,7 +22,6 @@ from typing import Dict, List, Optional, Set, Tuple, Type, Union
import torch import torch
import torch.distributed import torch.distributed
import torch_npu
from torch import nn from torch import nn
from vllm import envs from vllm import envs
from vllm.config import ParallelConfig, VllmConfig from vllm.config import ParallelConfig, VllmConfig
@@ -124,6 +123,8 @@ class NPUWorker(LocalOrDistributedWorkerBase):
# Torch profiler. Enabled and configured through env vars: # Torch profiler. Enabled and configured through env vars:
# VLLM_TORCH_PROFILER_DIR=/path/to/save/trace # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
if envs.VLLM_TORCH_PROFILER_DIR: if envs.VLLM_TORCH_PROFILER_DIR:
# lazy import so that torch_npu is not required for normal use.
import torch_npu
torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
logger.info("Profiling enabled. Traces will be saved to: %s", logger.info("Profiling enabled. Traces will be saved to: %s",
torch_profiler_trace_dir) torch_profiler_trace_dir)