wrap triton kernel with compat guard

Author: SigureMo

sonicmoe/functional/backward.py

@@ -20,6 +20,7 @@
 )
 from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton
 
+import paddle
 
 def _get_autotune_configs_for_db2_and_ds() -> list[triton.Config]:
     configs = []
@@ -28,6 +29,7 @@ def _get_autotune_configs_for_db2_and_ds() -> list[triton.Config]:
     return configs
 
 
+@paddle.use_compat_guard(enable=True, scope={"triton"})
 @triton.autotune(
     configs=_get_autotune_configs_for_db2_and_ds(),
     key=["H", "E"],
@@ -124,6 +126,7 @@ def _prune_triton_autotune_config(configs, nargs, **kw):
     return pruned_configs
 
 
+@paddle.use_compat_guard(enable=True, scope={"triton"})
 @triton.autotune(
     configs=_get_autotune_configs_for_db1(),
     key=["I", "E"],
@@ -168,6 +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"})
 @triton.jit
 def _colsum_smallN_kernel(
     y_ptr,  # *mut  T, shape [M]
@@ -484,6 +488,7 @@ def _token_broadcast_backward(
     )
 
 
+@paddle.use_compat_guard(enable=True, scope={"triton"})
 @triton.jit
 def _softmax_bwd_scatter_small_kernel(
     dlogits_ptr,
@@ -556,6 +561,7 @@ def _softmax_topk_bwd(
     )
 
 
+@paddle.use_compat_guard(enable=True, scope={"triton"})
 @triton.jit
 def _topk_bwd_scatter_small_kernel(
     dlogits_full_ptr,

sonicmoe/functional/forward.py

@@ -16,6 +16,8 @@
 from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton
 from .topk_softmax import TopK_Softmax
 
+import paddle
+
 
 @torch.library.custom_op(f"{LIBRARY_NAME}::_topk_fwd", mutates_args={"values", "indices"})
 def _topk_fwd(
@@ -202,6 +204,7 @@ def _router_forward(
     )
 
 
+@paddle.use_compat_guard(enable=True, scope={"triton"})
 @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,6 +21,8 @@
 from ..utils import get_next_power_of_2, get_powers_of_2
 from .tile_scheduler import SonicMoETileScheduler
 
+import paddle
+
 
 def last_even(a: int):
     return a if a % 2 == 0 else a - 1
@@ -438,6 +440,7 @@ def _prune_triton_autotune_config(configs, nargs, **kw):
         return pruned_configs
 
 
+@paddle.use_compat_guard(enable=True, scope={"triton"})
 @triton.autotune(
     configs=_get_triton_autotune_configs(),
     key=["H", "MAX_K", "w_is_None", "is_varlen_K"],