make triton optional if not available on system

Signed-off-by: cliu-us <[email protected]>

Author: chichun-charlie-liu

fms_mo/modules/linear.py

@@ -27,9 +27,6 @@
 import torch.nn.functional as F
 
 # Local
-from fms_mo.custom_ext_kernels.triton_kernels import (
-    tl_matmul_chunk_truncate as tl_matmul,
-)
 from fms_mo.custom_ext_kernels.utils import pack_vectorized
 from fms_mo.quant.quantizers import (
     HardPrune,
@@ -39,6 +36,13 @@
     get_weight_quantizer,
     mask_fc_kij,
 )
+from fms_mo.utils.import_utils import available_packages
+
+if available_packages["triton"]:
+    # Local
+    from fms_mo.custom_ext_kernels.triton_kernels import (
+        tl_matmul_chunk_truncate as tl_matmul,
+    )
 
 logger = logging.getLogger(__name__)
 
@@ -879,7 +883,9 @@ def from_torch_iW(cls, nnlin_iW, prec, a_cv, a_cvn, w_cv, zero_shift, **kwargs):
         qlinear_iW.nbits_w = 8
         qlinear_iW.acc_dtype = kwargs.get("acc_dtype", torch.float)
         qlinear_iW.usePTnativeQfunc = kwargs.get("use_PT_native_Qfunc", True)
-        qlinear_iW.use_int_kernel = kwargs.get("use_int_kernel", "triton")
+        qlinear_iW.use_int_kernel = kwargs.get(
+            "use_int_kernel", "triton" if available_packages["triton"] else False
+        )
         qlinear_iW.weight = nn.Parameter(
             nnlin_iW.weight.to(torch.int8), requires_grad=False
         )
@@ -1127,7 +1133,7 @@ def set_matmul_op(self):
                     chunk_size=self.chunk_size,
                 )
 
-            elif self.use_int_kernel == "cutlass":
+            elif self.use_int_kernel == "cutlass" and available_packages["cutlass"]:
                 # will use real imatmul written in cutlass
                 cutlass_ops_load_and_reg()
                 # Third Party

fms_mo/utils/import_utils.py

@@ -29,6 +29,7 @@
     "graphviz",
     "pygraphviz",
     "fms",
+    "triton",
 ]
 
 available_packages = {}

tests/triton_kernels/test_triton_mm.py

@@ -18,10 +18,14 @@
 import torch
 
 # Local
-from fms_mo.custom_ext_kernels.triton_kernels import (
-    tl_matmul_chunk_truncate as tl_matmul,
-)
 from fms_mo.modules.linear import LinearFPxAcc
+from fms_mo.utils.import_utils import available_packages
+
+if available_packages["triton"]:
+    # Local
+    from fms_mo.custom_ext_kernels.triton_kernels import (
+        tl_matmul_chunk_truncate as tl_matmul,
+    )
 
 
 @pytest.mark.parametrize("mkn", [64, 256, 1024, 4096])