Skip to content

GitLab

Unverified Commit 7101f4be authored by xiaoxi-wangfj's avatar xiaoxi-wangfj Committed by GitHub
Browse files

[PyTorch] Fix zero initialization in permute kernel for padded slots (#2026) 


Signed-off-by: default avatarxiaoxi-wangfj <690912414@qq.com>
parent 6c970612
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
transformer_engine/pytorch/triton/permutation.py
View file @ 7101f4be
......@@ -359,7 +359,7 @@ def _permute_kernel(
if prob == 0.0:
# for routing_map padding
# dst_row != -1 and prob == 0.0 means that this slot is padded
tl.store(output_ptr + output_off, 0, mask=mask)
tl.store(output_ptr + output_off, 0.0, mask=mask)
else:
tl.store(output_ptr + output_off, inp, mask=mask)
else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment