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

Enable intranode kernel tests with EP2 and EP4

Browse Source
This commit is contained in:
Chenggang Zhao 2025-03-03 15:01:02 +08:00
parent 6cc3497df8
commit c5b4040502
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tests/test_intranode.py
View File

@ -11,10 +11,11 @@ from utils import init_dist, bench, calc_diff, inplace_unique, per_token_cast_to
import test_low_latency
def test_main(num_sms: int, local_rank: int, num_local_ranks: int, num_ranks: int, rank: int, buffer: deep_ep.Buffer, group: dist.ProcessGroup):
def test_main(num_sms: int, local_rank: int, num_ranks: int, rank: int, buffer: deep_ep.Buffer, group: dist.ProcessGroup):
# Settings
# TODO: fix EP2/4/8 performance
num_tokens, hidden, num_topk, num_experts = 4096, 7168, 8, (256 // num_ranks) * num_ranks
assert num_experts % num_ranks == 0 and num_local_ranks == 8
assert num_experts % num_ranks == 0
if local_rank == 0:
print(f'[config] num_tokens={num_tokens}, hidden={hidden}, num_topk={num_topk}', flush=True)
@ -208,7 +209,7 @@ def test_loop(local_rank: int, num_local_ranks: int):
torch.manual_seed(rank)
for i in (24, ):
test_main(i, local_rank, num_local_ranks, num_ranks, rank, buffer, group)
test_main(i, local_rank, num_ranks, rank, buffer, group)
if local_rank == 0:
print()