diff --git a/backend/python/qwen-asr/backend.py b/backend/python/qwen-asr/backend.py index ea5e877f2d4f..a559ff5e678b 100644 --- a/backend/python/qwen-asr/backend.py +++ b/backend/python/qwen-asr/backend.py @@ -18,6 +18,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors from model_utils import resolve_model_reference +from device_utils import device_map_for, select_device @@ -95,13 +96,7 @@ def Health(self, request, context): return backend_pb2.Reply(message=bytes("OK", 'utf-8')) def LoadModel(self, request, context): - if torch.cuda.is_available(): - device = "cuda" - else: - device = "cpu" - mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available() - if mps_available: - device = "mps" + device = select_device(torch) if not torch.cuda.is_available() and request.CUDA: return backend_pb2.Result(success=False, message="CUDA is not available") @@ -123,7 +118,7 @@ def LoadModel(self, request, context): model_path, local_only = resolve_model_reference( request, "Qwen/Qwen3-ASR-1.7B" ) - default_dtype = torch.bfloat16 if self.device == "cuda" else torch.float32 + default_dtype = torch.bfloat16 if self.device in ("cuda", "xpu") else torch.float32 load_dtype = default_dtype if "torch_dtype" in self.options: d = str(self.options["torch_dtype"]).lower() @@ -145,12 +140,7 @@ def LoadModel(self, request, context): if attn_implementation is not None and isinstance(attn_implementation, str): attn_implementation = attn_implementation.strip() or None - if self.device == "mps": - device_map = None - elif self.device == "cuda": - device_map = "cuda:0" - else: - device_map = "cpu" + device_map = device_map_for(self.device) load_kwargs = dict( dtype=load_dtype, @@ -423,4 +413,4 @@ def signal_handler(sig, frame): parser = argparse.ArgumentParser(description="Run the gRPC server.") parser.add_argument("--addr", default="localhost:50051", help="The address to bind the server to.") args = parser.parse_args() - serve(args.addr) \ No newline at end of file + serve(args.addr) diff --git a/backend/python/qwen-asr/device_utils.py b/backend/python/qwen-asr/device_utils.py new file mode 100644 index 000000000000..0e1cb8005f11 --- /dev/null +++ b/backend/python/qwen-asr/device_utils.py @@ -0,0 +1,18 @@ +def select_device(torch_module): + mps = getattr(getattr(torch_module, "backends", None), "mps", None) + if mps is not None and mps.is_available(): + return "mps" + if torch_module.cuda.is_available(): + return "cuda" + xpu = getattr(torch_module, "xpu", None) + if xpu is not None and xpu.is_available(): + return "xpu" + return "cpu" + + +def device_map_for(device): + if device == "mps": + return None + if device in ("cuda", "xpu"): + return f"{device}:0" + return "cpu" diff --git a/backend/python/qwen-asr/device_utils_test.py b/backend/python/qwen-asr/device_utils_test.py new file mode 100644 index 000000000000..fde07f32b6e9 --- /dev/null +++ b/backend/python/qwen-asr/device_utils_test.py @@ -0,0 +1,58 @@ +import unittest + +from device_utils import device_map_for, select_device + + +class Availability: + def __init__(self, available): + self._available = available + + def is_available(self): + return self._available + + +class TorchStub: + def __init__(self, *, cuda=False, mps=False, xpu=False): + self.cuda = Availability(cuda) + self.backends = type("Backends", (), {"mps": Availability(mps)})() + self.xpu = Availability(xpu) + + +class SelectDeviceTest(unittest.TestCase): + def test_preserves_cuda_selection(self): + torch_module = TorchStub(cuda=True) + + self.assertEqual(select_device(torch_module), "cuda") + + def test_preserves_mps_selection(self): + torch_module = TorchStub(mps=True) + + self.assertEqual(select_device(torch_module), "mps") + + def test_selects_xpu_when_intel_gpu_is_available(self): + torch_module = TorchStub(xpu=True) + + self.assertEqual(select_device(torch_module), "xpu") + + def test_falls_back_to_cpu(self): + torch_module = TorchStub() + + self.assertEqual(select_device(torch_module), "cpu") + + +class DeviceMapTest(unittest.TestCase): + def test_preserves_cuda_model_placement(self): + self.assertEqual(device_map_for("cuda"), "cuda:0") + + def test_preserves_mps_model_placement(self): + self.assertIsNone(device_map_for("mps")) + + def test_places_the_model_on_the_first_xpu(self): + self.assertEqual(device_map_for("xpu"), "xpu:0") + + def test_preserves_cpu_model_placement(self): + self.assertEqual(device_map_for("cpu"), "cpu") + + +if __name__ == "__main__": + unittest.main()