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

Fix test combine args

Browse Source 
Signed-off-by: Hao Lin <linhaomails@gmail.com>
This commit is contained in:
Hao Lin 2025-04-11 18:21:09 +08:00
parent 8a0ca8e2ec
commit 23c54150ba
No known key found for this signature in database
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_internode.py
View File

@ -143,7 +143,7 @@ def test_main(num_sms: int, local_rank: int, num_local_ranks: int, num_ranks: in
if with_topk: if with_topk:
combine_args.update({'topk_weights': recv_topk_weights}) combine_args.update({'topk_weights': recv_topk_weights})
if previous_mode: if previous_mode:
dispatch_args.update({'previous_event': buffer.capture()}) combine_args.update({'previous_event': buffer.capture()})
combined_x, combined_topk_weights, event = buffer.combine(**combine_args) combined_x, combined_topk_weights, event = buffer.combine(**combine_args)
event.current_stream_wait() if async_mode else () event.current_stream_wait() if async_mode else ()
check_x = combined_x.float() / is_token_in_rank.sum(dim=1).unsqueeze(1) check_x = combined_x.float() / is_token_in_rank.sum(dim=1).unsqueeze(1)

2
tests/test_intranode.py
View File

@ -127,7 +127,7 @@ def test_main(num_sms: int, local_rank: int, num_ranks: int, rank: int, buffer:
if with_topk: if with_topk:
combine_args.update({'topk_weights': recv_topk_weights}) combine_args.update({'topk_weights': recv_topk_weights})
if previous_mode: if previous_mode:
dispatch_args.update({'previous_event': buffer.capture()}) combine_args.update({'previous_event': buffer.capture()})
combined_x, combined_topk_weights, event = buffer.combine(**combine_args) combined_x, combined_topk_weights, event = buffer.combine(**combine_args)
event.current_stream_wait() if async_mode else () event.current_stream_wait() if async_mode else ()
check_x = combined_x.float() / is_token_in_rank.sum(dim=1).unsqueeze(1) check_x = combined_x.float() / is_token_in_rank.sum(dim=1).unsqueeze(1)