Skip to content

GitLab

Unverified Commit 695b6347 authored by AlphaBaby's avatar AlphaBaby Committed by GitHub
Browse files

fix when --num-tokens == 1 (#356) 


Co-authored-by: default avatarfujianhao.fjh <fujianhao.fjh@antgroup.com>
parent 51071968
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
tests/test_low_latency.py
View file @ 695b6347
...@@ -77,6 +77,8 @@ def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int, ...@@ -77,6 +77,8 @@ def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int,
assert num_valid_tokens == (recv_layout_range & int_mask).sum().item(), f'{num_valid_tokens} != {recv_layout_range & int_mask}.sum().item()' assert num_valid_tokens == (recv_layout_range & int_mask).sum().item(), f'{num_valid_tokens} != {recv_layout_range & int_mask}.sum().item()'
assert num_valid_tokens == (all_topk_idx == expert_id).sum().item(), f'{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}' assert num_valid_tokens == (all_topk_idx == expert_id).sum().item(), f'{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}'
if num_valid_tokens == 0:
continue
# Check received data # Check received data
if current_x is x: if current_x is x:
recv_x = recv_x[:num_valid_tokens] recv_x = recv_x[:num_valid_tokens]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment