[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>

benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py

@@ -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(

benchmark/kernels/quantization/tuning_block_wise_kernel.py

@@ -31,7 +31,13 @@ from sglang.srt.layers.quantization.fp8_kernel import (
     _w8a8_block_fp8_matmul_unrolledx4,
 )
 from sglang.srt.layers.quantization.int8_kernel import _w8a8_block_int8_matmul
-from sglang.srt.utils import get_device_core_count, get_device_name, is_hip
+from sglang.srt.utils import (
+    get_device,
+    get_device_core_count,
+    get_device_count,
+    get_device_name,
+    is_hip,
+)
 
 _is_hip = is_hip()
 
@@ -221,18 +227,18 @@ def benchmark_config(
     def run():
         w8a8_block_matmul(A, B, As, Bs, block_size, config, out_dtype)
 
-    torch.cuda.synchronize()
+    torch.get_device_module().synchronize()
     # JIT complication & warmup
     for _ in range(5):
         run()
-    torch.cuda.synchronize()
+    torch.get_device_module().synchronize()
 
-    start_event = torch.cuda.Event(enable_timing=True)
-    end_event = torch.cuda.Event(enable_timing=True)
+    start_event = torch.get_device_module().Event(enable_timing=True)
+    end_event = torch.get_device_module().Event(enable_timing=True)
 
     latencies: List[float] = []
     for i in range(num_iters):
-        torch.cuda.synchronize()
+        torch.get_device_module().synchronize()
         start_event.record()
         run()
         end_event.record()
@@ -244,6 +250,7 @@ def benchmark_config(
 
 def tune(M, N, K, block_size, out_dtype, search_space, input_type):
     factor_for_scale = 1e-2
+    device = get_device()
 
     if input_type == "fp8":
         fp8_info = torch.finfo(
@@ -252,14 +259,14 @@ def tune(M, N, K, block_size, out_dtype, search_space, input_type):
         fp8_max, fp8_min = fp8_info.max, fp8_info.min
 
         A_fp32 = (
-            (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max
+            (torch.rand(M, K, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max
         )
         A = A_fp32.clamp(min=fp8_min, max=fp8_max).to(
             torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
         )
 
         B_fp32 = (
-            (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max
+            (torch.rand(N, K, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max
         )
         B = B_fp32.clamp(min=fp8_min, max=fp8_max).to(
             torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
@@ -269,12 +276,12 @@ def tune(M, N, K, block_size, out_dtype, search_space, input_type):
         int8_max, int8_min = int8_info.max, int8_info.min
 
         A_fp32 = (
-            (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * int8_max
+            (torch.rand(M, K, dtype=torch.float32, device=device) - 0.5) * 2 * int8_max
         )
         A = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8)
 
         B_fp32 = (
-            (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * int8_max
+            (torch.rand(N, K, dtype=torch.float32, device=device) - 0.5) * 2 * int8_max
         )
         B = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8)
 
@@ -282,9 +289,9 @@ def tune(M, N, K, block_size, out_dtype, search_space, input_type):
     n_tiles = (N + block_n - 1) // block_n
     k_tiles = (K + block_k - 1) // block_k
 
-    As = torch.rand(M, k_tiles, dtype=torch.float32, device="cuda") * factor_for_scale
+    As = torch.rand(M, k_tiles, dtype=torch.float32, device=device) * factor_for_scale
     Bs = (
-        torch.rand(n_tiles, k_tiles, dtype=torch.float32, device="cuda")
+        torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device)
         * factor_for_scale
     )
 
@@ -351,11 +358,6 @@ def save_configs(
             lock.release()
 
 
-def get_available_gpu_count():
-    """Get the number of available GPUs."""
-    return torch.cuda.device_count()
-
-
 def tune_on_gpu(args_dict):
     """Run tuning on a specific GPU."""
     gpu_id = args_dict["gpu_id"]
@@ -364,7 +366,7 @@ def tune_on_gpu(args_dict):
     args = args_dict["args"]
     lock = args_dict["lock"]
 
-    torch.cuda.set_device(gpu_id)
+    torch.get_device_module().set_device(gpu_id)
     print(f"Starting tuning on GPU {gpu_id} with batch sizes {batch_sizes}")
 
     block_n = args.block_n
@@ -415,12 +417,12 @@ def distribute_batch_sizes(batch_sizes, num_gpus):
 def main(args):
     print(args)
 
-    num_gpus = get_available_gpu_count()
+    num_gpus = get_device_count()
     if num_gpus == 0:
         raise RuntimeError("No GPU available for tuning")
     print(f"Found {num_gpus} GPUs for parallel tuning")
 
-    torch.cuda.init()
+    torch.get_device_module().init()
 
     if args.batch_size is None:
         batch_sizes = [