[Perf] Modify threshold for triggering allreduce norm fusion #12654
Conversation
Fridge003
requested review from
BBuf,
Edwardf0t1,
HaiShaw,
Ying1123,
ch-wan,
ispobock,
kushanam,
merrymercy and
zhyncs
as code owners
Summary of ChangesHello @Fridge003, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request refines the conditions for applying allreduce norm fusion, a performance optimization technique. By substantially lowering the maximum batch size threshold for this fusion, the change aims to optimize performance by restricting the fusion to smaller batch sizes, potentially improving efficiency or stability under specific operational loads. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request modifies the threshold for triggering all-reduce norm fusion, changing FUSE_ALLREDUCE_MAX_BATCH_SIZE from 2048 to 128. It also commendably replaces a magic number with this constant, improving code maintainability. The new threshold value is well-justified by an inline comment referencing a FlashInfer discussion. I have one suggestion to improve the clarity of the constant's name.
| from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant | ||
|
|
||
| FUSE_ALLREDUCE_MAX_BATCH_SIZE = 2048 | ||
| FUSE_ALLREDUCE_MAX_BATCH_SIZE = 128 |
There was a problem hiding this comment.
The constant name FUSE_ALLREDUCE_MAX_BATCH_SIZE appears to be misleading. Its usage throughout the file, such as the comparison with hidden_states.shape[0] on line 562, and the explanatory comment on line 556 which refers to "max token num", indicates that it represents a token limit, not a batch size limit. Renaming it to FUSE_ALLREDUCE_MAX_NUM_TOKENS would more accurately reflect its purpose and improve code clarity. If you make this change, please ensure all usages in the file are updated accordingly.
|
@yuan-luo @Qiaolin-Yu Can you have a look |
| and hasattr(layernorm, "forward_with_allreduce_fusion") | ||
| and get_global_server_args().enable_flashinfer_allreduce_fusion | ||
| and hidden_states.shape[0] <= 4096 | ||
| and hidden_states.shape[0] <= FUSE_ALLREDUCE_MAX_BATCH_SIZE |
There was a problem hiding this comment.
FLASHINFER_ALLREDUCE_FUSION_MAX_BATCH_SIZE might be better.
|
Have you tried on H20 or H800? The result might be slightly different. |
Not yet... I was doing benchmark on B200 and the result shows this change has negative effect. |
Motivation
Currently in DeepSeek fp4 and pure TP scenario,
flashinfer::trtllm_allreduce_fusionis applied to batches withbs <= FUSE_ALLREDUCE_MAX_BATCH_SIZE.However, with medium batch size like 2048(current threshold), we found that the cost for trtllm_allreduce_fusion might be much larger than all_reduce+norm with symmetric memory enabled. So we are lowering this threshold to a smaller value (like 128, as discussed in flashinfer-ai/flashinfer#1223 (comment))
Benchmarking and Profiling
DeepSeek fp4+TP 8 on B200
Profile (BS=1, ISL=2048, OSL=4):
Before this PR, the time of allreduce+norm fusion is 186us
After this PR, the time of allreduce kernel plus norm kernel with symmetric memory is 107us
Checklist