[Intel GPU] Enable DeepSeek R1 inference on XPU (#18461)
Signed-off-by: P V R K Jyothendra Varma <polisetty.v.r.k.jyothendra.varma@intel.com>
This commit is contained in:
@@ -32,9 +32,10 @@ from sglang.srt.server_args import (
|
||||
ServerArgs,
|
||||
set_global_server_args_for_scheduler,
|
||||
)
|
||||
from sglang.srt.utils import is_hip
|
||||
from sglang.srt.utils import get_device, is_hip, is_xpu
|
||||
|
||||
_is_hip = is_hip()
|
||||
_is_xpu = is_xpu()
|
||||
|
||||
|
||||
def benchmark_config(
|
||||
@@ -236,8 +237,8 @@ def benchmark_config(
|
||||
class BenchmarkWorker:
|
||||
|
||||
def __init__(self, seed: int, server_args: ServerArgs) -> None:
|
||||
torch.set_default_device("cuda")
|
||||
torch.cuda.manual_seed_all(0)
|
||||
torch.set_default_device(get_device())
|
||||
torch.get_device_module().manual_seed_all(0)
|
||||
self.seed = seed
|
||||
# Get the device ID to allocate tensors and kernels
|
||||
# on the respective GPU.
|
||||
@@ -330,7 +331,11 @@ class BenchmarkWorker:
|
||||
) -> Dict[str, int]:
|
||||
best_config = None
|
||||
best_time = float("inf")
|
||||
with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
|
||||
with (
|
||||
torch.get_device_module().device(self.device_id)
|
||||
if _is_xpu or _is_hip
|
||||
else nullcontext()
|
||||
):
|
||||
for config in tqdm(search_space):
|
||||
try:
|
||||
kernel_time = benchmark_config(
|
||||
|
||||
Reference in New Issue
Block a user