Update transformer_engine/pytorch/triton/permutation.py

Co-authored-by: Teddy Do <tdophung@nvidia.com>
Signed-off-by: xiaoxi-wangfj <690912414@qq.com>

Author: xiaoxi-wangfj

transformer_engine/pytorch/triton/permutation.py

@@ -454,7 +454,16 @@ def permute_with_mask_map(
                 (num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda"
             )
         else:
-            permuted_scale = torch.empty(
+    alloc = torch.zeros if pad_offsets is not None else torch.empty
+    output = alloc((num_out_tokens, hidden_size), dtype=inp.dtype, device="cuda")
+    permuted_probs = (
+        alloc((num_out_tokens,), dtype=probs.dtype, device="cuda") if probs is not None else None
+    )
+    permuted_scale = (
+        alloc((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda")
+        if scale is not None
+        else None
+    )
                 (num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda"
             )
     else: