Optimize the performance of the Combine Ant Moving function and the use of HCCL buffer

[oagniqgnat]

Motivation

The combine ant migration feature can effectively support long sequences when fewer HCCL buffers are used. Although the performance of the single-operator test with combine ant migration enabled is the same as that without it, the performance after integrating into the framework is poor, and the performance deteriorates significantly compared to before (2100 tps -> 1200 tps). This PR aims to improve the performance of combine ant migration while optimizing the use of HCCL buffers at the same time.

Modification

1. Optimize the performance of the Combine Ant Moving function using double buffer.
   Through profiling analysis, I found that the execution time of each stage within the operator does not significantly degrade. This is because the combine "ant moving" operation requires a full-card synchronization after each round of communication, and the model cannot guarantee that tokens will be evenly distributed to each expert, which results in significant performance overhead due to full-card synchronization. This PR addresses this issue by using double buffering in multiple rounds of communication, thereby overlapping the communication and full-card synchronization times, which greatly reduces the performance overhead caused by full-card synchronization. The only drawback of this approach is that it requires an additional buffer of size perRoundTokens * topK * hiddenSize. However, since perRoundTokens is typically set to 1024, this would only result in an additional 112MB buffer usage.

2. Optimize the performance of the Combine Ant Moving function by removing unnecessary PipeBarrier<PIPE_ALL>.

3. Optimize the HCCL normal operator by removing unnecessary notify_dispatch memory usage.
   The HCCL buffer usage of notify_dispatch is reduced from $204 MB \times 2$ to $102 MB \times 2$.

Benchmarking and Profiling

Our own single-operator benchmark

• disable ant moving: HCCL_BUFFSIZE=4000 python3 test_intranode.py --num-tokens=<bs>
• enable ant moving: HCCL_BUFFSIZE=900 DEEPEP_NORMAL_COMBINE_ENABLE_LONG_SEQ=1 DEEPEP_NORMAL_LONG_SEQ_ROUND=32 DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS=1024 python3 test_intranode.py --num-tokens=<bs>

bs	disable GB/s	enable (before optimization) GB/s	enable (after optimization) GB/s
4096	101.59	98.38	109.01
8192	101.30	96.72	113.56
16384	99.30	97.55	113.00
32768	-	98.24	114.52

Qwen 235B prefill

disable token/s	enable (before optimization) token/s	enable (after optimization) token/s
2040	1200	2106

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments (/tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci) or contact authorized users to do so.
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!