Guard triton.set_allocator with hasattr for non-NVIDIA backends

Author: AriesMTR

src/flag_blas/runtime/backend/_iluvatar/__init__.py

@@ -1,4 +1,4 @@
-from flag_gems.runtime.backend.backend_utils import VendorInfoBase
+from flag_blas.runtime.backend.backend_utils import VendorInfoBase
 
 vendor_info = VendorInfoBase(
     vendor_name="iluvatar", device_name="cuda", device_query_cmd="ixsmi"

src/flag_blas/runtime/backend/_nvidia/__init__.py

@@ -13,7 +13,8 @@ def _alloc_fn(size, alignment, stream):
     return torch.empty(size, device="cuda", dtype=torch.int8)
 
 
-triton.set_allocator(_alloc_fn)
+if hasattr(triton, "set_allocator"):
+    triton.set_allocator(_alloc_fn)
 
 
 __all__ = ["*"]