Merge pull request bytedance#108 from bytedance/fix_micro_perf

[micro_perf] fix log and sqrt in gpu backend.
hliuca · Sep 27, 2024 · bfb03e6 · bfb03e6
2 parents c3d23ad + 6b1d531
commit bfb03e6
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/byte_micro_perf/backends/GPU/backend_gpu.py b/byte_micro_perf/backends/GPU/backend_gpu.py
@@ -88,11 +88,7 @@ def initialize_ccl(self, rank, world_size):
         self.setup_2d_group()
         return True
 
-    def log(self):
-        self.op = LogOp()
 
-    def sqrt(self):
-        self.op = SqrtOp()
 
     def setup_2d_group(self):
         # get rank and set device
@@ -247,6 +243,11 @@ def swiglu(self):
     def cast(self):
         self.op = module_store.CastOp()
 
+    def log(self):
+        self.op = module_store.LogOp()
+
+    def sqrt(self):
+        self.op = module_store.SqrtOp()
 
     # binary ops
     def add(self):