mindstudio-probe 1.0.1__py3-none-any.whl → 1.0.4__py3-none-any.whl

msprobe/pytorch/free_benchmark/result_handlers/check_handler.py

@@ -1,39 +1,39 @@
-from typing import Any
-
-from msprobe.pytorch.free_benchmark import logger
-from msprobe.pytorch.free_benchmark.common.enums import DeviceType
-from msprobe.pytorch.free_benchmark.common.params import DataParams, make_unequal_row
-from msprobe.pytorch.free_benchmark.common.utils import Tools
-from msprobe.pytorch.free_benchmark.compare.single_benchmark import SingleCompare
-from msprobe.pytorch.free_benchmark.result_handlers.base_handler import FuzzHandler
-
-
-class CheckerHandler(FuzzHandler):
-    def other_compare(self, data_params: DataParams) -> bool:
-        is_consistent = SingleCompare().compare_seq(
-                    data_params.original_result, data_params.perturbed_result
-                )
-        if not is_consistent:
-            self.unequal_rows.append(
-                make_unequal_row(data_params, self.params)
-            )
-
-    def get_threshold(self, dtype):
-        return self._get_default_threshold(dtype)
-
-    def handle(self, data_params: DataParams) -> Any:
-        if isinstance(data_params.perturbed_result, bool) or not Tools.is_float_tensor(
-            data_params.perturbed_result
-        ):
-            return data_params.original_result
-        try:
-            if self.params.fuzz_device == DeviceType.NPU:
-                self.cmp_output_npu(data_params)
-            else:
-                self.other_compare(data_params)
-        except Exception as e:
-            logger.warning_on_rank_0(
-                f"[msprobe] Free Benchmark: For {self.params.api_name}, "
-                f"when campare the result exception raise {e}"
-            )
-        return data_params.original_result
+from typing import Any
+
+from msprobe.pytorch.free_benchmark import logger
+from msprobe.pytorch.free_benchmark.common.enums import DeviceType
+from msprobe.pytorch.free_benchmark.common.params import DataParams, make_unequal_row
+from msprobe.pytorch.free_benchmark.common.utils import Tools
+from msprobe.pytorch.free_benchmark.compare.single_benchmark import SingleCompare
+from msprobe.pytorch.free_benchmark.result_handlers.base_handler import FuzzHandler
+
+
+class CheckerHandler(FuzzHandler):
+    def other_compare(self, data_params: DataParams) -> bool:
+        is_consistent = SingleCompare().compare_seq(
+                    data_params.original_result, data_params.perturbed_result
+                )
+        if not is_consistent:
+            self.unequal_rows.append(
+                make_unequal_row(data_params, self.params)
+            )
+
+    def get_threshold(self, dtype):
+        return self._get_default_threshold(dtype)
+
+    def handle(self, data_params: DataParams) -> Any:
+        if isinstance(data_params.perturbed_result, bool) or not Tools.is_float_tensor(
+            data_params.perturbed_result
+        ):
+            return data_params.original_result
+        try:
+            if self.params.fuzz_device == DeviceType.NPU:
+                self.cmp_output_npu(data_params)
+            else:
+                self.other_compare(data_params)
+        except Exception as e:
+            logger.warning_on_rank_0(
+                f"[msprobe] Free Benchmark: For {self.params.api_name}, "
+                f"when campare the result exception raise {e}"
+            )
+        return data_params.original_result

msprobe/pytorch/free_benchmark/result_handlers/fix_handler.py

@@ -1,24 +1,24 @@
-from typing import Any
-
-from msprobe.pytorch.free_benchmark.common.params import DataParams
-from msprobe.pytorch.free_benchmark.common.utils import Tools
-from msprobe.pytorch.free_benchmark.result_handlers.base_handler import FuzzHandler
-from msprobe.pytorch.free_benchmark import logger
-
-
-class FixHandler(FuzzHandler):
-
-    def get_threshold(self, dtype):
-        return self._get_default_threshold(dtype)
-
-    def handle(self, data_params: DataParams) -> Any:
-        try:
-            return Tools.convert_fuzz_output_to_origin(
-                data_params.original_result, data_params.perturbed_result
-            )
-        except Exception as e:
-            logger.warning_on_rank_0(
-                f"[msprobe] Free Benchmark: For {self.params.api_name} "
-                f"Fix output failed. "
-            )
+from typing import Any
+
+from msprobe.pytorch.free_benchmark.common.params import DataParams
+from msprobe.pytorch.free_benchmark.common.utils import Tools
+from msprobe.pytorch.free_benchmark.result_handlers.base_handler import FuzzHandler
+from msprobe.pytorch.free_benchmark import logger
+
+
+class FixHandler(FuzzHandler):
+
+    def get_threshold(self, dtype):
+        return self._get_default_threshold(dtype)
+
+    def handle(self, data_params: DataParams) -> Any:
+        try:
+            return Tools.convert_fuzz_output_to_origin(
+                data_params.original_result, data_params.perturbed_result
+            )
+        except Exception as e:
+            logger.warning_on_rank_0(
+                f"[msprobe] Free Benchmark: For {self.params.api_name} "
+                f"Fix output failed. "
+            )
         return data_params.original_result

msprobe/pytorch/free_benchmark/result_handlers/handler_factory.py

@@ -1,31 +1,30 @@
-from msprobe.pytorch.free_benchmark import FreeBenchmarkException
-from msprobe.pytorch.free_benchmark.common.constant import PreheatConfig
-from msprobe.pytorch.free_benchmark.common.enums import HandlerType
-from msprobe.pytorch.free_benchmark.common.params import HandlerParams
-from msprobe.pytorch.free_benchmark.result_handlers.check_handler import CheckerHandler
-from msprobe.pytorch.free_benchmark.result_handlers.preheat_handler import PreheatHandler
-from msprobe.pytorch.free_benchmark.result_handlers.fix_handler import FixHandler
-
-
-class FuzzHandlerFactory:
-
-    result_handlers = {
-        HandlerType.CHECK: CheckerHandler,
-        HandlerType.FIX: FixHandler,
-        HandlerType.PREHEAT: PreheatHandler,
-    }
-
-    @staticmethod
-    def create(params: HandlerParams):
-        if_preheat = params.preheat_config.get(PreheatConfig.IF_PREHEAT)
-        if not if_preheat:
-            handler = FuzzHandlerFactory.result_handlers.get(params.handler_type)
-        else:
-            handler = FuzzHandlerFactory.result_handlers.get(HandlerType.PREHEAT)
-            # TODO
-        if not handler:
-            raise FreeBenchmarkException(
-                FreeBenchmarkException.UnsupportedType,
-                f"无标杆工具支持 [ {HandlerType.CHECK}、{HandlerType.FIX}] 形式",
-            )
-        return handler(params)
+from msprobe.pytorch.free_benchmark import FreeBenchmarkException
+from msprobe.pytorch.free_benchmark.common.constant import PreheatConfig
+from msprobe.pytorch.free_benchmark.common.enums import HandlerType
+from msprobe.pytorch.free_benchmark.common.params import HandlerParams
+from msprobe.pytorch.free_benchmark.result_handlers.check_handler import CheckerHandler
+from msprobe.pytorch.free_benchmark.result_handlers.preheat_handler import PreheatHandler
+from msprobe.pytorch.free_benchmark.result_handlers.fix_handler import FixHandler
+
+
+class FuzzHandlerFactory:
+
+    result_handlers = {
+        HandlerType.CHECK: CheckerHandler,
+        HandlerType.FIX: FixHandler,
+        HandlerType.PREHEAT: PreheatHandler,
+    }
+
+    @staticmethod
+    def create(params: HandlerParams):
+        if_preheat = params.preheat_config.get(PreheatConfig.IF_PREHEAT)
+        if not if_preheat:
+            handler = FuzzHandlerFactory.result_handlers.get(params.handler_type)
+        else:
+            handler = FuzzHandlerFactory.result_handlers.get(HandlerType.PREHEAT)
+        if not handler:
+            raise FreeBenchmarkException(
+                FreeBenchmarkException.UnsupportedType,
+                f"无标杆工具支持 [ {HandlerType.CHECK}、{HandlerType.FIX}] 形式",
+            )
+        return handler(params)

msprobe/pytorch/free_benchmark/result_handlers/preheat_handler.py

@@ -1,170 +1,170 @@
-import math
-from typing import Any
-
-from msprobe.pytorch.free_benchmark import logger
-from msprobe.pytorch.free_benchmark.common.constant import ThresholdConfig
-from msprobe.pytorch.free_benchmark.common.counter import preheat_counter
-from msprobe.pytorch.free_benchmark.common.enums import DeviceType
-from msprobe.pytorch.free_benchmark.common.params import DataParams, HandlerParams
-from msprobe.pytorch.free_benchmark.common.utils import Tools
-from msprobe.pytorch.free_benchmark.compare.single_benchmark import SingleCompare
-from msprobe.pytorch.free_benchmark.result_handlers.base_handler import FuzzHandler
-
-
-class PreheatHandler(FuzzHandler):
-
-    def __init__(self, params: HandlerParams) -> None:
-        super().__init__(params)
-        self.pure_name = Tools.get_pure_api_name(self.params.api_name)
-
-    def get_threshold(self, dtype):
-        return preheat_counter.get_api_thd(self.pure_name, dtype)
-
-    def compare_npu_and_cpu(self, data_params: DataParams):
-        args = Tools.convert_device_and_dtype(
-            data_params.args, DeviceType.CPU, change_dtype=True
-        )
-        kwargs = Tools.convert_device_and_dtype(
-            data_params.kwargs, DeviceType.CPU, change_dtype=True
-        )
-        cpu_result = data_params.origin_func(*args, **kwargs)
-        return SingleCompare().compare_seq(data_params.original_result, cpu_result)
-
-    def preheat(self, max_fuzz_ratio, cpu_consistent, first_dtype):
-        # 存储当前step所有输出比值和对应npu\cpu比对结果
-        preheat_counter.update_preheat_record(
-            self.pure_name,
-            first_dtype,
-            (max_fuzz_ratio, cpu_consistent),
-        )
-        if self._need_adjust_threshold():
-            self._adjust_threshold()
-
-    def handle(self, data_params: DataParams) -> Any:
-
-        if isinstance(data_params.perturbed_result, bool) or not Tools.is_float_tensor(
-                data_params.perturbed_result
-        ):
-            return data_params.original_result
-
-        if self.params.step == 0:
-            preheat_counter.add_one_step_used_api(self.pure_name)
-            return data_params.original_result
-
-        # 如果当前api,step需要预热
-        npu_consistent, max_fuzz_ratio = self.cmp_output_npu(data_params)
-        data_params.is_consistent = npu_consistent
-
-        preheat_counter.check_step(self.params.step)
-
-        if self.params.preheat_config.get("preheat_step") <= self.params.step:
-            return data_params.original_result
-
-        if not data_params.grad_unequal_flag:
-            data_params.grad_unequal_flag = True
-            data_params.is_consistent = False
-            return data_params.original_result
-        preheat_counter.add_api_called_time(self.pure_name)
-
-        if not self._is_take_a_sample():
-            return data_params.original_result
-
-        cpu_consistent = True
-        try:
-            cpu_consistent = self.compare_npu_and_cpu(data_params)
-        except Exception as e:
-            logger.warning_on_rank_0(
-                f"[msprobe] Free Benchmark: For {self.params.api_name}, "
-                f"when campare to cpu exception raise {e}"
-            )
-        try:
-            first_dtype = Tools.get_first_tensor_dtype(data_params.original_result)
-        except RuntimeError:
-            logger.warning_on_rank_0(
-                f"[msprobe] Free Benchmark: For {self.params.api_name}, "
-                f"the output sequence does not contain tensors."
-            )
-        if preheat_counter.get_api_preheat(self.pure_name, str(first_dtype)):
-            self.preheat(max_fuzz_ratio, cpu_consistent, first_dtype)
-
-        return data_params.original_result
-
-    def _is_take_a_sample(self) -> bool:
-        need_sample_set = self._get_need_sample_set()
-        curr_called_seq = preheat_counter.get_api_called_time(self.pure_name)
-        res = curr_called_seq in need_sample_set
-        if res:
-            total_count = preheat_counter.get_one_step_used_api(self.pure_name)
-            logger.info_on_rank_0(
-                f"[msprobe] Free benchmark: preheat sample in step{self.params.step}"
-                f"api_name {self.params.api_name}, "
-                f"curr_called_seq: {curr_called_seq}/{total_count}"
-            )
-            preheat_counter.add_api_sample_time(self.pure_name)
-        return res
-
-    def _get_sample_count_per_step(self) -> set:
-        """
-        每一个step中应该采集的样本数
-        """
-        total_count = preheat_counter.get_one_step_used_api(self.pure_name)
-        preheat_step = self.params.preheat_config.get("preheat_step")
-        max_sample = self.params.preheat_config.get("max_sample")
-        return min(math.ceil(total_count / preheat_step), max_sample)
-
-    def _get_need_sample_set(self):
-        """
-        需要采集的api集合
-        """
-        # 每一步样本数
-        total_count = preheat_counter.get_one_step_used_api(self.pure_name)
-        sample_count_per_step = self._get_sample_count_per_step()
-        need_sample_set = set()
-        prehead_step = self.params.preheat_config.get("preheat_step")
-        for i in range(1, sample_count_per_step + 1):
-            count = (prehead_step * (i - 1) + self.params.step) % total_count
-            if count == 0:
-                count = total_count
-            need_sample_set.add(count)
-        return need_sample_set
-
-    def _need_adjust_threshold(self) -> bool:
-        sample_count_per_step = self._get_sample_count_per_step()
-        sampled_time = preheat_counter.get_api_sample_time(self.pure_name)
-        res = sampled_time >= sample_count_per_step
-        return res
-
-    def _adjust_threshold_for_dtype(self, dtype_str, compare_result):
-        con_ratio = [ratio for ratio, is_consistent in compare_result if is_consistent]
-        incon_ratio = [
-            ratio for ratio, is_consistent in compare_result if not is_consistent
-        ]
-        old_thd = preheat_counter.get_api_thd(self.pure_name, dtype_str)
-        new_thd = old_thd
-        # 正例负例都存在
-        if con_ratio and incon_ratio:
-            if min(incon_ratio) > max(con_ratio):
-                new_thd = min(min(incon_ratio), old_thd)
-                preheat_counter.set_api_preheat(self.pure_name, dtype_str, is_preheat=False)
-        elif con_ratio:
-            # 存在漏报
-            if max(con_ratio) > old_thd:
-                new_thd = 1 + ((old_thd - 1) * ThresholdConfig.API_THD_STEP)
-            else:
-                new_thd = 1 + ((old_thd - 1) / ThresholdConfig.API_THD_STEP)
-        else:
-            new_thd = min(min(incon_ratio), old_thd)
-            preheat_counter.set_api_preheat(self.pure_name, dtype_str, is_preheat=False)
-        return new_thd
-
-    def _adjust_threshold(self):
-        for dtype_str, compare_result in preheat_counter.preheat_record[
-            self.pure_name
-        ].items():
-            new_thd = self._adjust_threshold_for_dtype(dtype_str, compare_result)
-            threshold = self._get_default_threshold(
-                preheat_counter.dtype_map.get(dtype_str)
-            )
-            preheat_counter.update_api_thd(
-                self.pure_name, dtype_str, new_thd, threshold
-            )
+import math
+from typing import Any
+
+from msprobe.pytorch.free_benchmark import logger
+from msprobe.pytorch.free_benchmark.common.constant import ThresholdConfig
+from msprobe.pytorch.free_benchmark.common.counter import preheat_counter
+from msprobe.pytorch.free_benchmark.common.enums import DeviceType
+from msprobe.pytorch.free_benchmark.common.params import DataParams, HandlerParams
+from msprobe.pytorch.free_benchmark.common.utils import Tools
+from msprobe.pytorch.free_benchmark.compare.single_benchmark import SingleCompare
+from msprobe.pytorch.free_benchmark.result_handlers.base_handler import FuzzHandler
+
+
+class PreheatHandler(FuzzHandler):
+
+    def __init__(self, params: HandlerParams) -> None:
+        super().__init__(params)
+        self.pure_name = Tools.get_pure_api_name(self.params.api_name)
+
+    def get_threshold(self, dtype):
+        return preheat_counter.get_api_thd(self.pure_name, dtype)
+
+    def compare_npu_and_cpu(self, data_params: DataParams):
+        args = Tools.convert_device_and_dtype(
+            data_params.args, DeviceType.CPU, change_dtype=True
+        )
+        kwargs = Tools.convert_device_and_dtype(
+            data_params.kwargs, DeviceType.CPU, change_dtype=True
+        )
+        cpu_result = data_params.origin_func(*args, **kwargs)
+        return SingleCompare().compare_seq(data_params.original_result, cpu_result)
+
+    def preheat(self, max_fuzz_ratio, cpu_consistent, first_dtype):
+        # 存储当前step所有输出比值和对应npu\cpu比对结果
+        preheat_counter.update_preheat_record(
+            self.pure_name,
+            first_dtype,
+            (max_fuzz_ratio, cpu_consistent),
+        )
+        if self._need_adjust_threshold():
+            self._adjust_threshold()
+
+    def handle(self, data_params: DataParams) -> Any:
+
+        if isinstance(data_params.perturbed_result, bool) or not Tools.is_float_tensor(
+                data_params.perturbed_result
+        ):
+            return data_params.original_result
+
+        if self.params.step == 0:
+            preheat_counter.add_one_step_used_api(self.pure_name)
+            return data_params.original_result
+
+        # 如果当前api,step需要预热
+        npu_consistent, max_fuzz_ratio = self.cmp_output_npu(data_params)
+        data_params.is_consistent = npu_consistent
+
+        preheat_counter.check_step(self.params.step)
+
+        if self.params.preheat_config.get("preheat_step") <= self.params.step:
+            return data_params.original_result
+
+        if not data_params.grad_unequal_flag:
+            data_params.grad_unequal_flag = True
+            data_params.is_consistent = False
+            return data_params.original_result
+        preheat_counter.add_api_called_time(self.pure_name)
+
+        if not self._is_take_a_sample():
+            return data_params.original_result
+
+        cpu_consistent = True
+        try:
+            cpu_consistent = self.compare_npu_and_cpu(data_params)
+        except Exception as e:
+            logger.warning_on_rank_0(
+                f"[msprobe] Free Benchmark: For {self.params.api_name}, "
+                f"when campare to cpu exception raise {e}"
+            )
+        try:
+            first_dtype = Tools.get_first_tensor_dtype(data_params.original_result)
+        except RuntimeError:
+            logger.warning_on_rank_0(
+                f"[msprobe] Free Benchmark: For {self.params.api_name}, "
+                f"the output sequence does not contain tensors."
+            )
+        if preheat_counter.get_api_preheat(self.pure_name, str(first_dtype)):
+            self.preheat(max_fuzz_ratio, cpu_consistent, first_dtype)
+
+        return data_params.original_result
+
+    def _is_take_a_sample(self) -> bool:
+        need_sample_set = self._get_need_sample_set()
+        curr_called_seq = preheat_counter.get_api_called_time(self.pure_name)
+        res = curr_called_seq in need_sample_set
+        if res:
+            total_count = preheat_counter.get_one_step_used_api(self.pure_name)
+            logger.info_on_rank_0(
+                f"[msprobe] Free benchmark: preheat sample in step{self.params.step}"
+                f"api_name {self.params.api_name}, "
+                f"curr_called_seq: {curr_called_seq}/{total_count}"
+            )
+            preheat_counter.add_api_sample_time(self.pure_name)
+        return res
+
+    def _get_sample_count_per_step(self) -> set:
+        """
+        每一个step中应该采集的样本数
+        """
+        total_count = preheat_counter.get_one_step_used_api(self.pure_name)
+        preheat_step = self.params.preheat_config.get("preheat_step")
+        max_sample = self.params.preheat_config.get("max_sample")
+        return min(math.ceil(total_count / preheat_step), max_sample)
+
+    def _get_need_sample_set(self):
+        """
+        需要采集的api集合
+        """
+        # 每一步样本数
+        total_count = preheat_counter.get_one_step_used_api(self.pure_name)
+        sample_count_per_step = self._get_sample_count_per_step()
+        need_sample_set = set()
+        prehead_step = self.params.preheat_config.get("preheat_step")
+        for i in range(1, sample_count_per_step + 1):
+            count = (prehead_step * (i - 1) + self.params.step) % total_count
+            if count == 0:
+                count = total_count
+            need_sample_set.add(count)
+        return need_sample_set
+
+    def _need_adjust_threshold(self) -> bool:
+        sample_count_per_step = self._get_sample_count_per_step()
+        sampled_time = preheat_counter.get_api_sample_time(self.pure_name)
+        res = sampled_time >= sample_count_per_step
+        return res
+
+    def _adjust_threshold_for_dtype(self, dtype_str, compare_result):
+        con_ratio = [ratio for ratio, is_consistent in compare_result if is_consistent]
+        incon_ratio = [
+            ratio for ratio, is_consistent in compare_result if not is_consistent
+        ]
+        old_thd = preheat_counter.get_api_thd(self.pure_name, dtype_str)
+        new_thd = old_thd
+        # 正例负例都存在
+        if con_ratio and incon_ratio:
+            if min(incon_ratio) > max(con_ratio):
+                new_thd = min(min(incon_ratio), old_thd)
+                preheat_counter.set_api_preheat(self.pure_name, dtype_str, is_preheat=False)
+        elif con_ratio:
+            # 存在漏报
+            if max(con_ratio) > old_thd:
+                new_thd = 1 + ((old_thd - 1) * ThresholdConfig.API_THD_STEP)
+            else:
+                new_thd = 1 + ((old_thd - 1) / ThresholdConfig.API_THD_STEP)
+        else:
+            new_thd = min(min(incon_ratio), old_thd)
+            preheat_counter.set_api_preheat(self.pure_name, dtype_str, is_preheat=False)
+        return new_thd
+
+    def _adjust_threshold(self):
+        for dtype_str, compare_result in preheat_counter.preheat_record[
+            self.pure_name
+        ].items():
+            new_thd = self._adjust_threshold_for_dtype(dtype_str, compare_result)
+            threshold = self._get_default_threshold(
+                preheat_counter.dtype_map.get(dtype_str)
+            )
+            preheat_counter.update_api_thd(
+                self.pure_name, dtype_str, new_thd, threshold
+            )

msprobe/pytorch/function_factory.py

@@ -0,0 +1,76 @@
+from msprobe.pytorch.common.utils import logger
+from msprobe.pytorch.bench_functions.apply_adam_w import npu_apply_adam_w
+from msprobe.pytorch.bench_functions.confusion_transpose import npu_confusion_transpose, \
+    npu_confusion_transpose_backward
+from msprobe.pytorch.bench_functions.fast_gelu import npu_fast_gelu, npu_fast_gelu_backward
+from msprobe.pytorch.bench_functions.layer_norm_eval import npu_layer_norm_eval
+from msprobe.pytorch.bench_functions.linear import npu_linear, npu_linear_backward
+from msprobe.pytorch.bench_functions.matmul_backward import matmul_backward
+from msprobe.pytorch.bench_functions.npu_fusion_attention import npu_fusion_attention, npu_fusion_attention_grad, \
+    gpu_fusion_attention
+from msprobe.pytorch.bench_functions.rms_norm import npu_rms_norm, npu_rms_norm_backward
+from msprobe.pytorch.bench_functions.rotary_mul import npu_rotary_mul, npu_rotary_mul_backward
+from msprobe.pytorch.bench_functions.scaled_mask_softmax import npu_scaled_masked_softmax, \
+    npu_scaled_masked_softmax_backward
+from msprobe.pytorch.bench_functions.swiglu import npu_swiglu, npu_swiglu_backward, swish_grad, swish
+
+
+class Register(dict):
+    def __init__(self, *args, **kwargs):
+        super(Register, self).__init__(*args, **kwargs)
+        self._dict = {}
+
+    def __call__(self, target_func_list):
+        for target in target_func_list:
+            self.register(target)
+        return
+
+    def __setitem__(self, key, value):
+        self._dict[key] = value
+
+    def __getitem__(self, key):
+        return self._dict[key]
+
+    def __contains__(self, key):
+        return key in self._dict
+
+    def __str__(self):
+        return str(self._dict)
+
+    def keys(self):
+        return self._dict.keys()
+
+    def values(self):
+        return self._dict.values()
+
+    def items(self):
+        return self._dict.items()
+
+    def register(self, target):
+
+        def add_register_item(key, value):
+            if key in self._dict:
+                logger.warning(f"{value.__name__} has been registered before, so we will overriden it.")
+            self[key] = value
+            return value
+
+        if callable(target):
+            return add_register_item(target.__name__, target)
+        else:
+            raise Exception(f"The func {target} is not callable.")
+
+
+# register for npu custom bench functions
+npu_custom_functions = Register()
+npu_custom_functions([
+    npu_apply_adam_w, npu_confusion_transpose, npu_fast_gelu, npu_layer_norm_eval, npu_linear, npu_fusion_attention,
+    npu_rms_norm, npu_rotary_mul, npu_scaled_masked_softmax, npu_swiglu, gpu_fusion_attention
+])
+
+# register for npu custom backward bench functions
+npu_custom_grad_functions = Register()
+npu_custom_grad_functions([
+    npu_confusion_transpose_backward, npu_fast_gelu_backward, npu_linear_backward, matmul_backward,
+    npu_fusion_attention_grad, npu_rms_norm_backward, npu_rotary_mul_backward, npu_scaled_masked_softmax_backward,
+    npu_swiglu_backward
+])