Correct setting of permuted_probs_stride_token, unpermuted_probs_stride_token and unpermuted_probs_stride_expert in unpermutation

Author: mingxu1067

transformer_engine/jax/triton/permutation.py

@@ -354,9 +354,9 @@ def grid_fn(meta):
         output_stride_hidden,
         merging_probs_stride_token,
         merging_probs_stride_expert,
-        1,
-        0,
-        0,
+        permuted_probs_stride_token,
+        unpermuted_probs_stride_token,
+        unpermuted_probs_stride_expert,
         kernel=_unpermute_kernel,
         out_shape=out_shape,
         grid=grid_fn,