updata moe_utils.py

megatron/core/transformer/moe/moe_utils.py

@@ -19,7 +19,7 @@ try:
 except ImportError:
     HAVE_TE = False
 
-
+@torch.compile(options={"triton.cudagraphs": True, "triton.cudagraph_trees": False, "triton.cudagraph_support_input_mutation":True})
 def switch_load_balancing_loss_func(
     probs: torch.Tensor,
     tokens_per_expert: torch.Tensor,
@@ -217,7 +217,7 @@ class MoEAuxLossAutoScaler(torch.autograd.Function):
         """
         MoEAuxLossAutoScaler.main_loss_backward_scale = scale
 
-
+@torch.compile(mode='max-autotune-no-cudagraphs')
 def permute(
     tokens,
     routing_map,
@@ -278,7 +278,7 @@ def permute(
 
     return permuted_input, sorted_indices
 
-
+@torch.compile(mode='max-autotune-no-cudagraphs')
 def unpermute(
     permuted_tokens: torch.Tensor,
     sorted_indices: torch.Tensor,