[PyTorch] Fix garbage initialized permuted_scale (NVIDIA#2547)

xiaoxi-wangfj · tdophung · web-flow · commit c988548f72bb · 2026-01-02T15:23:29.000-08:00
Signed-off-by: xiaoxi-wangfj &lt;690912414@qq.com&gt;
Co-authored-by: Teddy Do &lt;tdophung@nvidia.com&gt;
diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py
@@ -165,7 +165,7 @@ def permute_with_mask_map(
         alloc((num_out_tokens,), dtype=probs.dtype, device="cuda") if probs is not None else None
     )
     permuted_scale = (
-        torch.empty((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda")
+        alloc((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda")
         if scale is not None
         else None
     )

Original file line number	Diff line number	Diff line change
`@@ -165,7 +165,7 @@ def permute_with_mask_map(`
`165`	`165`	`alloc((num_out_tokens,), dtype=probs.dtype, device="cuda") if probs is not None else None`
`166`	`166`	`)`
`167`	`167`	`permuted_scale = (`
`168`		`- torch.empty((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda")`
	`168`	`+ alloc((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda")`
`169`	`169`	`if scale is not None`
`170`	`170`	`else None`
`171`	`171`	`)`