fb57a6162f
Signed-off-by: KexiongYu <yukexiong1@huawei.com> Signed-off-by: betta18 <jiangmengyu1@huawei.com> Signed-off-by: muziyuhui666 <111362884+muziyuhui666@users.noreply.github.com> Signed-off-by: jiangmengyu18 <56633611+jiangmengyu18@users.noreply.github.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: jiangmengyu18 <56633611+jiangmengyu18@users.noreply.github.com> Co-authored-by: KexiongYu <yukexiong1@huawei.com> Co-authored-by: betta18 <jiangmengyu1@huawei.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import torch.nn as nn
|
|
|
|
from vllm_omni.utils.platform_utils import detect_device_type, is_npu, is_rocm
|
|
|
|
|
|
class CustomOp(nn.Module):
|
|
"""
|
|
Base class for custom ops.
|
|
Dispatches the forward method to the appropriate backend.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.is_cuda = detect_device_type() == "cuda"
|
|
self._forward_method = self.dispatch_forward()
|
|
|
|
def dispatch_forward(self) -> Callable:
|
|
if is_rocm():
|
|
return self.forward_hip
|
|
elif self.is_cuda:
|
|
return self.forward_cuda
|
|
elif is_npu():
|
|
return self.forward_npu
|
|
else:
|
|
return self.forward_native
|
|
|
|
def forward(self, *args, **kwargs) -> Any:
|
|
return self._forward_method(*args, **kwargs)
|
|
|
|
def forward_native(self, *args, **kwargs):
|
|
"""PyTorch-native implementation of the forward method.
|
|
This method is optional. If implemented, it can be used with compilers
|
|
such as torch.compile or PyTorch XLA. Also, it can be used for testing
|
|
purposes.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def forward_cuda(self, *args, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
def forward_npu(self, *args, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
def forward_hip(self, *args, **kwargs):
|
|
# By default, we assume that HIP ops are compatible with CUDA ops.
|
|
return self.forward_cuda(*args, **kwargs)
|