From 545e420ce614e6cd7605d3284f6994cbb06c74ea Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 22 Jan 2025 16:56:01 -0800
Subject: [PATCH] move TritonOps

---
 unsloth/__init__.py          | 35 +++++++----------------------------
 unsloth/matmul_perf_model.py | 17 ++++++++++++++++-
 2 files changed, 23 insertions(+), 29 deletions(-)

diff --git a/unsloth/__init__.py b/unsloth/__init__.py
index 74a0adce57..d1c9ab9ef6 100644
--- a/unsloth/__init__.py
+++ b/unsloth/__init__.py
@@ -143,33 +143,10 @@ if Version(triton.__version__) >= Version("3.0.0"):
     except: pass
 else: from triton.common.build import libcuda_dirs
 
-def fix_triton_ops():
-    # Check if triton.ops exists
-    try:
-        import triton.ops
-    except:
-        # Triton 3.2 removed triton.ops
-        from .matmul_perf_model import (
-            early_config_prune   as _early_config_prune,
-            estimate_matmul_time as _estimate_matmul_time,
-        )
-        class PerfOps:
-            def __init__(self): return
-            @staticmethod
-            def early_config_prune(*args, **kwargs):
-                return _early_config_prune(*args, **kwargs)
-            @staticmethod
-            def estimate_matmul_time(*args, **kwargs):
-                return _estimate_matmul_time(*args, **kwargs)
-        pass
-        class TritonOps:
-            __slots__ = "matmul_perf_model",
-            def __init__(self): self.matmul_perf_model = PerfOps()
-        pass
-        triton.ops = TritonOps()
-    pass
-pass
-fix_triton_ops()
+# Triton 3.2 removed triton.ops, so we shall fix it!
+from .matmul_perf_model import TritonOps
+try: import triton.ops
+except: triton.ops = TritonOps()
 
 # Try loading bitsandbytes and triton
 import bitsandbytes as bnb
@@ -210,7 +187,9 @@ except:
         else: from triton.common.build import libcuda_dirs
         cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
         libcuda_dirs()
-        fix_triton_ops()
+        # Triton 3.2 removed triton.ops, so we shall fix it!
+        try: import triton.ops
+        except: triton.ops = TritonOps()
     except:
         warnings.warn(
             "Unsloth: CUDA is not linked properly.\n"\
diff --git a/unsloth/matmul_perf_model.py b/unsloth/matmul_perf_model.py
index 53b59d808e..6a86c29bd4 100644
--- a/unsloth/matmul_perf_model.py
+++ b/unsloth/matmul_perf_model.py
@@ -208,4 +208,19 @@ def early_config_prune(configs, named_args, **kwargs):
             random_config = v[0][0]
             random_config.num_stages = 2
             pruned_configs.append(random_config)
-    return pruned_configs
\ No newline at end of file
+    return pruned_configs
+
+class PerfOps:
+    def __init__(self): return
+    @staticmethod
+    def early_config_prune(*args, **kwargs):
+        return _early_config_prune(*args, **kwargs)
+    @staticmethod
+    def estimate_matmul_time(*args, **kwargs):
+        return _estimate_matmul_time(*args, **kwargs)
+pass
+
+class TritonOps:
+    __slots__ = "matmul_perf_model",
+    def __init__(self): self.matmul_perf_model = PerfOps()
+pass