[Perf] Reduce memory usage by splitting tokens in fused_experts and avoiding unused tensor #833
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
ApsarasX
force-pushed
the
wengang/memory-optimization
branch
6 times, most recently
from
27ec0bd to
1516602
Compare
ApsarasX
changed the title
ApsarasX
force-pushed
the
wengang/memory-optimization
branch
from
1516602 to
276a46d
Compare
ApsarasX
force-pushed
the
wengang/memory-optimization
branch
from
276a46d to
f8a0e4b
Compare
ApsarasX
force-pushed
the
wengang/memory-optimization
branch
from
f8a0e4b to
c323f91
Compare
Signed-off-by: ApsarasX <apsarax@outlook.com>
Signed-off-by: ApsarasX <apsarax@outlook.com>
ApsarasX
force-pushed
the
wengang/memory-optimization
branch
from
c323f91 to
ab355ab
Compare
wangxiyuan
approved these changes
May 19, 2025
| topk_ids_list = topk_ids.split(VLLM_FUSED_EXPERTS_SEQ_SPLIT_LENGTH) | ||
| final_hidden_states_list = [] | ||
| for i in range(len(x_list)): | ||
| final_hidden_states = fused_experts( |
There was a problem hiding this comment.
Do you have some assessment on the performance of this PR? This might bring some substantial performance regression, for this change may lead to repeat weight access from HBM
There was a problem hiding this comment.
maybe the env VLLM_FUSED_EXPERTS_SEQ_SPLIT_LENGTH can be changed to something like ENABLE_FUSED_EXPOERS_SEQ_SPLIT to keep the backward capibility
|
The optimizations in this PR are specific to the A2 machine and conflict with the alltoall-based optimizations you're using on A3 machine. Therefore, I will temporarily close this PR for now and reopen it later if needed. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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?
splitting tokens in fused_experts
When
--max-model-len=32768 onDeepSeek-R1-W8A8, thefused_expertsfunction consumes about 5.75GB of memory. By splitting it into multiple executions, the memory consumption of thefused_expertsfunction can be reduced to 1.2GB, thereby increasing the available KVCache.The disadvantage of this solution is that when the number of prompt tokens sent by the user exceeds VLLM_FUSED_EXPERTS_SEQ_SPLIT_LENGTH, an additional
concatoperator overhead will be added.However, considering that the user's request in most scenarios will be less than
8192, we believe that this overhead is acceptable.avoiding unused tensor
self.inputs_embedsin NPUModelRunner V1 will always be generated, but it will only be used in multi-modal situations, so I changed its generation conditions to reduce memory usage.Does this PR introduce any user-facing change?
No
How was this patch tested?