Improve Prefill Performance by Removing Redundant Padding and Optimizing Alltoall Communication #948
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What this PR does / why we need it?
This PR improves Prefill performance by making two key optimizations:
all_to_all_single
call followed by threeall_to_all
calls. This has been refactored to use threeall_to_all_single
calls instead, with a fixed communication buffer to eliminate an extra communication step. This change not only simplifies the communication pattern but also leverages the better performance ofall_to_all_single
.While there might be minor precision trade-offs, the choice of the coefficient
2
is an empirically sound value that maintains accuracy even when expert ID distribution is imbalanced.In testing with DeepSeek-V3, the model was able to handle 3584-token inputs with significantly improved Prefill throughput and no regression in dialog quality.
Does this PR introduce any user-facing change?
No, this PR does not introduce any user-facing changes.
How was this patch tested?