diff --git a/.github/workflows/pr-test-npu.yml b/.github/workflows/pr-test-npu.yml index ad02fcee5..4ee7e3f2d 100644 --- a/.github/workflows/pr-test-npu.yml +++ b/.github/workflows/pr-test-npu.yml @@ -74,6 +74,8 @@ jobs: HCCL_BUFFSIZE: 1913 run: | python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py + python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py --num-tokens=1 + python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py --num-tokens=2 - name: Run test base fused deep moe timeout-minutes: 10 @@ -168,6 +170,8 @@ jobs: HCCL_BUFFSIZE: 1913 run: | python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py + python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py --num-tokens=1 + python3 $GITHUB_WORKSPACE/tests/python/deepep/test_low_latency.py --num-tokens=2 - name: Run test base fused deep moe timeout-minutes: 10 diff --git a/tests/python/deepep/test_low_latency.py b/tests/python/deepep/test_low_latency.py index 1c359afc9..6bebad3f5 100644 --- a/tests/python/deepep/test_low_latency.py +++ b/tests/python/deepep/test_low_latency.py @@ -74,110 +74,112 @@ def test( cumulative_local_expert_recv_stats = torch.zeros( (num_local_experts,), dtype=torch.int, device="npu" ) - dispatch_use_fp8 = True - packed_recv_x, packed_recv_count, handle, event, hook = buffer.low_latency_dispatch( - x, - topk_idx, - num_tokens, - num_experts, - use_fp8=dispatch_use_fp8, - round_scale=False, - use_ue8m0=False, - cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats, - async_finish=not return_recv_hook, - return_recv_hook=return_recv_hook, - ) - simulated_gemm_x = ( - per_token_cast_back(*packed_recv_x) if dispatch_use_fp8 else packed_recv_x - ) - - all_topk_idx = torch.empty( - (num_ranks, num_tokens, num_topk), dtype=topk_idx.dtype, device="npu" - ) - dist.all_gather_into_tensor(all_topk_idx, topk_idx, group=group) - - for i in range(num_local_experts if do_check else 0): - expert_id = rank * num_local_experts + i - temp = num_tokens / num_local_experts - recv_count = packed_recv_count[i] - recv_x = ( - per_token_cast_back( - packed_recv_x[0][int(i * temp) : int((i + 1) * temp)], - packed_recv_x[1][int(i * temp) : int((i + 1) * temp)], + for dispatch_use_fp8 in (True, False): + packed_recv_x, packed_recv_count, handle, event, hook = ( + buffer.low_latency_dispatch( + x, + topk_idx, + num_tokens, + num_experts, + use_fp8=dispatch_use_fp8, + round_scale=False, + use_ue8m0=False, + cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats, + async_finish=not return_recv_hook, + return_recv_hook=return_recv_hook, ) - if dispatch_use_fp8 - else packed_recv_x[int(i * temp) : int((i + 1) * temp)] ) - if i == 0: - recv_layout_range = handle[1][(i + 1) * num_ranks - 1] - else: - recv_layout_range = ( - handle[1][(i + 1) * num_ranks - 1] - handle[1][i * num_ranks - 1] - ) + simulated_gemm_x = ( + per_token_cast_back(*packed_recv_x) if dispatch_use_fp8 else packed_recv_x + ) - # Check expert indices - int_mask = (2**32) - 1 - num_valid_tokens = recv_count.item() - assert ( - num_valid_tokens == (recv_layout_range & int_mask).item() - ), f"{num_valid_tokens} != {recv_layout_range & int_mask}.item()" - assert ( - num_valid_tokens == (all_topk_idx == expert_id).sum().item() - ), f"{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}" + all_topk_idx = torch.empty( + (num_ranks, num_tokens, num_topk), dtype=topk_idx.dtype, device="npu" + ) + dist.all_gather_into_tensor(all_topk_idx, topk_idx, group=group) - if num_valid_tokens == 0: - continue - # Check received data - recv_x = recv_x[:num_valid_tokens] - recv_x_amin = recv_x[:, :-128].amin(dim=-1) - assert torch.equal(recv_x_amin, recv_x[:, :-128].amax(dim=-1)) - if dispatch_use_fp8: - hash_value ^= hash_tensor( - packed_recv_x[0][int(i * temp) : int(i * temp + num_valid_tokens)] - ) - hash_value ^= hash_tensor( - packed_recv_x[1][int(i * temp) : int(i * temp + num_valid_tokens)] - ) - else: - hash_value ^= hash_tensor( - packed_recv_x[int(i * temp) : int(i * temp + num_valid_tokens)] + for i in range(num_local_experts if do_check else 0): + expert_id = rank * num_local_experts + i + temp = num_tokens / num_local_experts + recv_count = packed_recv_count[i] + recv_x = ( + per_token_cast_back( + packed_recv_x[0][int(i * temp) : int((i + 1) * temp)], + packed_recv_x[1][int(i * temp) : int((i + 1) * temp)], + ) + if dispatch_use_fp8 + else packed_recv_x[int(i * temp) : int((i + 1) * temp)] ) + if i == 0: + recv_layout_range = handle[1][(i + 1) * num_ranks - 1] + else: + recv_layout_range = ( + handle[1][(i + 1) * num_ranks - 1] - handle[1][i * num_ranks - 1] + ) - # Check combine correctness - ( - src_info, - layout_range, - num_max_dispatch_tokens_per_rank, - hidden, - num_experts, - packed_recv_count, - ) = handle + # Check expert indices + int_mask = (2**32) - 1 + num_valid_tokens = recv_count.item() + assert ( + num_valid_tokens == (recv_layout_range & int_mask).item() + ), f"{num_valid_tokens} != {recv_layout_range & int_mask}.item()" + assert ( + num_valid_tokens == (all_topk_idx == expert_id).sum().item() + ), f"{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}" - out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="npu") - combined_x, event, hook = buffer.low_latency_combine( - simulated_gemm_x, - topk_idx, - topk_weights, - handle, - async_finish=not return_recv_hook, - zero_copy=False, - return_recv_hook=return_recv_hook, - out=out, - ) + if num_valid_tokens == 0: + continue + # Check received data + recv_x = recv_x[:num_valid_tokens] + recv_x_amin = recv_x[:, :-128].amin(dim=-1) + assert torch.equal(recv_x_amin, recv_x[:, :-128].amax(dim=-1)) + if dispatch_use_fp8: + hash_value ^= hash_tensor( + packed_recv_x[0][int(i * temp) : int(i * temp + num_valid_tokens)] + ) + hash_value ^= hash_tensor( + packed_recv_x[1][int(i * temp) : int(i * temp + num_valid_tokens)] + ) + else: + hash_value ^= hash_tensor( + packed_recv_x[int(i * temp) : int(i * temp + num_valid_tokens)] + ) - if do_check: - diff = calc_diff( - x * topk_weights.masked_fill(topk_idx == -1, 0).sum(dim=1).view(-1, 1), - combined_x, + # Check combine correctness + ( + src_info, + layout_range, + num_max_dispatch_tokens_per_rank, + hidden, + num_experts, + packed_recv_count, + ) = handle + + out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="npu") + combined_x, event, hook = buffer.low_latency_combine( + simulated_gemm_x, + topk_idx, + topk_weights, + handle, + async_finish=not return_recv_hook, + zero_copy=False, + return_recv_hook=return_recv_hook, + out=out, ) - assert torch.isnan(combined_x).sum().item() == 0 - if dispatch_use_fp8: - assert diff < 1e-4, f"Error: {diff=}" - else: - assert diff < 1e-5, f"Error: {diff=}" - hash_value ^= hash_tensor(combined_x) - print(f"rank {rank} PASSED") + if do_check: + diff = calc_diff( + x * topk_weights.masked_fill(topk_idx == -1, 0).sum(dim=1).view(-1, 1), + combined_x, + ) + assert torch.isnan(combined_x).sum().item() == 0 + if dispatch_use_fp8: + assert diff < 1e-4, f"Error: {diff=}" + else: + assert diff < 1e-5, f"Error: {diff=}" + hash_value ^= hash_tensor(combined_x) + + print(f"rank {rank} PASSED") # noinspection PyShadowingNames def test_func(zero_copy: bool, return_recv_hook: bool):