wrap triton kernel with getitem

Author: SigureMo

sonicmoe/functional/backward.py

@@ -20,7 +20,7 @@
 )
 from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton
 
-import paddle
+from ..triton_utils import wrap_triton_kernel
 
 def _get_autotune_configs_for_db2_and_ds() -> list[triton.Config]:
     configs = []
@@ -29,7 +29,7 @@ def _get_autotune_configs_for_db2_and_ds() -> list[triton.Config]:
     return configs
 
 
-@paddle.use_compat_guard(enable=True, scope={"triton"}, silent=True)
+@wrap_triton_kernel
 @triton.autotune(
     configs=_get_autotune_configs_for_db2_and_ds(),
     key=["H", "E"],
@@ -126,7 +126,7 @@ def _prune_triton_autotune_config(configs, nargs, **kw):
     return pruned_configs
 
 
-@paddle.use_compat_guard(enable=True, scope={"triton"}, silent=True)
+@wrap_triton_kernel
 @triton.autotune(
     configs=_get_autotune_configs_for_db1(),
     key=["I", "E"],
@@ -171,7 +171,7 @@ def db1_kernel(
         tl.store(db1_ptr + Eidx * I + i_offsets, db1_acc, mask=i_mask)
 
 
-@paddle.use_compat_guard(enable=True, scope={"triton"}, silent=True)
+@wrap_triton_kernel
 @triton.jit
 def _colsum_smallN_kernel(
     y_ptr,  # *mut  T, shape [M]
@@ -488,7 +488,7 @@ def _token_broadcast_backward(
     )
 
 
-@paddle.use_compat_guard(enable=True, scope={"triton"}, silent=True)
+@wrap_triton_kernel
 @triton.jit
 def _softmax_bwd_scatter_small_kernel(
     dlogits_ptr,
@@ -561,7 +561,7 @@ def _softmax_topk_bwd(
     )
 
 
-@paddle.use_compat_guard(enable=True, scope={"triton"}, silent=True)
+@wrap_triton_kernel
 @triton.jit
 def _topk_bwd_scatter_small_kernel(
     dlogits_full_ptr,

sonicmoe/functional/forward.py

@@ -16,7 +16,7 @@
 from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton
 from .topk_softmax import TopK_Softmax
 
-import paddle
+from ..triton_utils import wrap_triton_kernel
 
 
 @torch.library.custom_op(f"{LIBRARY_NAME}::_topk_fwd", mutates_args={"values", "indices"})
@@ -204,7 +204,7 @@ def _router_forward(
     )
 
 
-@paddle.use_compat_guard(enable=True, scope={"triton"}, silent=True)
+@wrap_triton_kernel
 @triton.jit
 def _softmax_fwd_small_kernel(
     logits_ptr, stride_lm: tl.constexpr, stride_ln: tl.constexpr, K: tl.constexpr, BLOCK_K: tl.constexpr

sonicmoe/functional/reduction_over_k_gather.py

@@ -21,7 +21,7 @@
 from ..utils import get_next_power_of_2, get_powers_of_2
 from .tile_scheduler import SonicMoETileScheduler
 
-import paddle
+from ..triton_utils import wrap_triton_kernel
 
 
 def last_even(a: int):
@@ -440,7 +440,7 @@ def _prune_triton_autotune_config(configs, nargs, **kw):
         return pruned_configs
 
 
-@paddle.use_compat_guard(enable=True, scope={"triton"}, silent=True)
+@wrap_triton_kernel
 @triton.autotune(
     configs=_get_triton_autotune_configs(),
     key=["H", "MAX_K", "w_is_None", "is_varlen_K"],

sonicmoe/triton_utils.py

@@ -0,0 +1,10 @@
+import paddle
+
+def wrap_triton_kernel(triton_kernel):
+    class WrappedTritonKernel:
+        def __init__(self, kernel):
+            self.kernel = kernel
+
+        def __getitem__(self, index):
+            return paddle.use_compat_guard(enable=True, scope={"triton"})(self.kernel[index])
+    return WrappedTritonKernel(triton_kernel)