[Inductor] Update set_driver_to_gpu code to avoid backend re-initialization with new Triton (#147621)

anmyachev · pytorchmergebot · commit 97557b983348 · 2025-02-25T00:04:54.000Z
Pull Request resolved: #147621 Approved by: https://github.com/jansel
diff --git a/torch/_inductor/runtime/triton_helpers.py b/torch/_inductor/runtime/triton_helpers.py
@@ -27,7 +27,13 @@ def set_driver_to_gpu():
     driver = triton.runtime.driver
     for name, backend in triton.backends.backends.items():
         if backend.driver.is_active() and name != "cpu":
-            if isinstance(driver.active, backend.driver):
+            # After https://github.com/triton-lang/triton/commit/b844d519bc5e86edf00fe6b3c6c2d1badcd509a4,
+            # `driver.active` can be of `LazyProxy` type and the sign of this - `_obj` attribute.
+            if (
+                isinstance(driver.active, backend.driver)
+                or hasattr(driver.active, "_obj")
+                and isinstance(driver.active._obj, backend.driver)
+            ):
                 # Don't re-initialize backend if it is already active
                 return
             driver.set_active(backend.driver())
