Skip to content

GitLab

"...git@developer.sourcefind.cn:kecinstone/2024-pra-vllm.git" did not exist on "75c0ca9d432199d336e0bc6ae854aa1d35ba35ce"
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( ...@@ -359,7 +359,7 @@ def _permute_kernel(
if prob == 0.0: if prob == 0.0:
# for routing_map padding # for routing_map padding
# dst_row != -1 and prob == 0.0 means that this slot is padded # 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: else:
tl.store(output_ptr + output_off, inp, mask=mask) tl.store(output_ptr + output_off, inp, mask=mask)
else: 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