We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
1 parent c5a4f44 commit 67af0b1Copy full SHA for 67af0b1
1 file changed
fms_mo/custom_ext_kernels/triton_kernels.py
@@ -15,11 +15,23 @@
15
"""This file contains external kernels for FP and INT8 matmul written in triton."""
16
17
# Third Party
18
-from triton.language.extra import libdevice
19
import torch
+
20
+# First Party
21
+from fms_mo.utils.import_utils import available_packages
22
23
+# Assume any calls to the file are requesting triton
24
+if not available_packages["triton"]:
25
+ raise ImportError(
26
+ "triton python package is not avaialble, please check your installation."
27
+ )
28
29
+# Third Party
30
+from triton.language.extra import libdevice
31
import triton
32
import triton.language as tl
33
34
35
DTYPE_I8 = [torch.int8]
36
DTYPE_F8 = [torch.float8_e4m3fn, torch.float8_e5m2]
37
DTYPE_8BIT = DTYPE_I8 + DTYPE_F8
0 commit comments