DeepSeek/DeepEP
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepEP synced 2025-06-26 18:28:11 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Improve EP2/4 performance

Browse Source
This commit is contained in:
Chenggang Zhao
2025-03-04 15:30:19 +08:00
parent 55cdd9a64f
commit 1553fc42bf
3 changed files with 85 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
deep_ep/buffer.py
View File

@@ -160,12 +160,11 @@ class Buffer:
Returns:
config: the recommended config.
"""
# Intranode
if num_ranks <= 8:
return Config(Buffer.num_sms, 6, 256, 6, 128)
# Internode
config_map = {
2: Config(Buffer.num_sms, 16, 256, 6, 128),
4: Config(Buffer.num_sms, 16, 256, 6, 128),
8: Config(Buffer.num_sms, 6, 256, 6, 128),
16: Config(Buffer.num_sms, 16, 288, 20, 128),
24: Config(Buffer.num_sms, 8, 288, 32, 128),
32: Config(Buffer.num_sms, 8, 288, 32, 128),
@@ -188,12 +187,11 @@ class Buffer:
Returns:
config: the recommended config.
"""
# Intranode
if num_ranks <= 8:
return Config(Buffer.num_sms, 6, 256, 6, 128)
# Internode
config_map = {
2: Config(Buffer.num_sms, 6, 256, 6, 128),
4: Config(Buffer.num_sms, 6, 256, 6, 128),
8: Config(Buffer.num_sms, 6, 256, 6, 128),
16: Config(Buffer.num_sms, 2, 288, 28, 128),
24: Config(Buffer.num_sms, 1, 288, 20, 128),
32: Config(Buffer.num_sms, 1, 288, 20, 128),