Try to use triton.language.extra.libdevice when possible (#3196)

Summary:
Pull Request resolved: #3196

X-link: pytorch/pytorch#136997

X-link: facebookresearch/generative-recommenders#90

X-link: facebookresearch/FBGEMM#294

In view of triton-lang/triton#3825 we should try to use `triton.language.extra.libdevice` instead of `triton.language.extra.cuda.libdevice`.

Reviewed By: bertmaher, karthik-man

Differential Revision: D63583965

fbshipit-source-id: d32f35f7524d45c1e7c95c095144ad27f16eaa5a

fbgemm_gpu/experimental/gen_ai/test/kv_cache/rope_padded.py

@@ -38,8 +38,12 @@
         # pyre-fixme[21]: Could not find name `pow` in `triton.language.math`.
         from triton.language.math import pow
     except ImportError:
-        # @manual=//triton:triton
-        from triton.language.extra.cuda.libdevice import pow
+        try:
+            # @manual=//triton:triton
+            from triton.language.extra.libdevice import pow
+        except ImportError:
+            # @manual=//triton:triton
+            from triton.language.extra.cuda.libdevice import pow
 
 
 _INTERNAL_DTYPE_MAP: Dict[str, int] = {"": 0, "f32": 1, "f64": 2}