Skip to content

Fix the device error when using ray as vllm-acend backend #884

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix the device error when using ray as vllm-acend backend #884

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion tests/ops/test_rotary_embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,12 @@
import torch
import torch.nn as nn

import vllm_ascend.platform # noqa: F401
from vllm_ascend.utils import try_register_lib

try_register_lib(
"vllm_ascend.vllm_ascend_C",
exc_info=
"Warning: Failed to register custom ops, all custom ops will be disabled.")

# Only Neox style true scenario is supported for now
IS_NEOX_STYLE = [True]
Expand Down
9 changes: 7 additions & 2 deletions vllm_ascend/ops/rotary_embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,11 +22,16 @@
from vllm.model_executor.layers.rotary_embedding import (
DeepseekScalingRotaryEmbedding, RotaryEmbedding)

from vllm_ascend.platform import CUSTOM_OP_ENABLED
from vllm_ascend.utils import try_register_lib


def custom_rotary_embedding_enabled(query, neox_style, head_size):
return query.dtype == torch.float16 and neox_style and head_size % 32 == 0 and CUSTOM_OP_ENABLED
try_register_lib(
"vllm_ascend.vllm_ascend_C",
exc_info=
"Warning: Failed to register custom ops, all custom ops will be disabled."
)
return query.dtype == torch.float16 and neox_style and head_size % 32 == 0


def rope_forward_oot(
Expand Down
16 changes: 0 additions & 16 deletions vllm_ascend/platform.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,6 @@
# This file is a part of the vllm-ascend project.
#

import logging
import os
from typing import TYPE_CHECKING, Optional, Tuple

import torch
Expand All @@ -27,18 +25,6 @@

from vllm_ascend.utils import ASCEND_QUATIZATION_METHOD, update_aclgraph_sizes

CUSTOM_OP_ENABLED = False
try:
# register custom ops into torch_library here
import vllm_ascend.vllm_ascend_C # type: ignore # noqa: F401

except ImportError:
logging.warning(
"Warning: Failed to register custom ops, all custom ops will be disabled"
)
else:
CUSTOM_OP_ENABLED = True

if TYPE_CHECKING:
from vllm.config import ModelConfig, VllmConfig
from vllm.utils import FlexibleArgumentParser
Expand All @@ -47,8 +33,6 @@
VllmConfig = None
FlexibleArgumentParser = None

os.environ["RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES"] = "1"


class NPUPlatform(Platform):

Expand Down
4 changes: 3 additions & 1 deletion vllm_ascend/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@
ASCEND_QUATIZATION_METHOD = "ascend"


def try_register_lib(lib_name: str, lib_info: str = ""):
def try_register_lib(lib_name: str, lib_info: str = "", exc_info: str = ""):
import importlib
import importlib.util
try:
Expand All @@ -51,6 +51,8 @@ def try_register_lib(lib_name: str, lib_info: str = ""):
if lib_info:
logger.info(lib_info)
except Exception:
if exc_info:
logger.info(exc_info)
pass


Expand Down
Loading