e7e1a7dc05
### What this PR does / why we need it?
#5051 only implement a basic framework for model runner v2, but there
are still some bugs for e2e functionality, this PR aim to enable basic
functionality.
model runner v2 plans:
https://github.com/vllm-project/vllm-ascend/issues/5208
- vLLM version: release/v0.13.0
- vLLM main:
ad32e3e19c
---------
Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
21 lines
615 B
Python
21 lines
615 B
Python
from contextlib import contextmanager
|
|
|
|
import torch
|
|
|
|
|
|
@contextmanager
|
|
def torch_cuda_wrapper():
|
|
try:
|
|
torch.cuda.Event = torch.npu.Event
|
|
torch.cuda.Stream = torch.npu.Stream
|
|
torch.cuda.stream = torch.npu.stream
|
|
torch.cuda.default_stream = torch.npu.default_stream
|
|
torch.cuda.current_stream = torch.npu.current_stream
|
|
torch.cuda.graph_pool_handle = torch.npu.graph_pool_handle
|
|
torch.cuda.CUDAGraph = torch.npu.NPUGraph
|
|
torch.cuda.graph = torch.npu.graph
|
|
torch.cuda.synchronize = torch.npu.synchronize
|
|
yield
|
|
finally:
|
|
pass
|