From f2aa00cae6e1d1b1a5a085604455c5c80bc5673a Mon Sep 17 00:00:00 2001 From: zhuyy Date: Tue, 6 Jan 2026 19:43:23 +0800 Subject: [PATCH 1/5] =?UTF-8?q?fix=20little=20batchsize=20and=20int8?= =?UTF-8?q?=E9=87=8F=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .github/workflows/pr-test-npu.yml | 10 +- tests/python/deepep/test_low_latency.py | 190 ++++++++++++------------ 2 files changed, 104 insertions(+), 96 deletions(-) diff --git a/.github/workflows/pr-test-npu.yml b/.github/workflows/pr-test-npu.yml index ad02fcee5..6a19f9621 100644 --- a/.github/workflows/pr-test-npu.yml +++ b/.github/workflows/pr-test-npu.yml @@ -57,7 +57,9 @@ jobs: env: HCCL_BUFFSIZE: 2300 run: | - python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py + python3 $GITHUB_WORKSPACE/tests/python/deepep/ + python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=1 + python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 - name: Run test multi-round intranode timeout-minutes: 10 @@ -74,6 +76,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 @@ -161,6 +165,8 @@ jobs: HCCL_BUFFSIZE: 2300 run: | python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py + python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=1 + python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 - name: Run test low latency timeout-minutes: 10 @@ -168,6 +174,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..bd2869b83 100644 --- a/tests/python/deepep/test_low_latency.py +++ b/tests/python/deepep/test_low_latency.py @@ -74,110 +74,110 @@ 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)], - ) - if dispatch_use_fp8 - else packed_recv_x[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, + ) + simulated_gemm_x = ( + per_token_cast_back(*packed_recv_x) if dispatch_use_fp8 else packed_recv_x ) - 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 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): From 262ec5ff749d0dd8a73bd677d6fcd7943077034c Mon Sep 17 00:00:00 2001 From: zhuyy Date: Tue, 6 Jan 2026 20:21:59 +0800 Subject: [PATCH 2/5] fix some pipeline errors --- .github/workflows/pr-test-npu.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/workflows/pr-test-npu.yml b/.github/workflows/pr-test-npu.yml index 6a19f9621..c63300074 100644 --- a/.github/workflows/pr-test-npu.yml +++ b/.github/workflows/pr-test-npu.yml @@ -57,8 +57,7 @@ jobs: env: HCCL_BUFFSIZE: 2300 run: | - python3 $GITHUB_WORKSPACE/tests/python/deepep/ - python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=1 + python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 - name: Run test multi-round intranode @@ -165,7 +164,6 @@ jobs: HCCL_BUFFSIZE: 2300 run: | python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py - python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=1 python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 - name: Run test low latency From 3a5589c2d57f1d4b031ef310e3b3ec31aac368ba Mon Sep 17 00:00:00 2001 From: zhuyy Date: Wed, 7 Jan 2026 15:40:46 +0800 Subject: [PATCH 3/5] delete little bs command. Please enter the commit message for your changes. Lines starting --- .github/workflows/pr-test-npu.yml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pr-test-npu.yml b/.github/workflows/pr-test-npu.yml index c63300074..a0c73ec51 100644 --- a/.github/workflows/pr-test-npu.yml +++ b/.github/workflows/pr-test-npu.yml @@ -58,7 +58,8 @@ jobs: HCCL_BUFFSIZE: 2300 run: | python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py - python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 + # python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=1 + # python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 - name: Run test multi-round intranode timeout-minutes: 10 @@ -164,7 +165,8 @@ jobs: HCCL_BUFFSIZE: 2300 run: | python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py - python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 + # python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=1 + # python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 - name: Run test low latency timeout-minutes: 10 From 718a7d380a77e5c997d1267e699e85b5fcb08d7c Mon Sep 17 00:00:00 2001 From: zhuyy Date: Thu, 8 Jan 2026 10:45:13 +0800 Subject: [PATCH 4/5] fix grammatical error --- .github/workflows/pr-test-npu.yml | 6 ++---- tests/python/deepep/test_low_latency.py | 24 +++++++++++++----------- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/.github/workflows/pr-test-npu.yml b/.github/workflows/pr-test-npu.yml index a0c73ec51..c63300074 100644 --- a/.github/workflows/pr-test-npu.yml +++ b/.github/workflows/pr-test-npu.yml @@ -58,8 +58,7 @@ jobs: HCCL_BUFFSIZE: 2300 run: | python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py - # python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=1 - # python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 + python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 - name: Run test multi-round intranode timeout-minutes: 10 @@ -165,8 +164,7 @@ jobs: HCCL_BUFFSIZE: 2300 run: | python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py - # python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=1 - # python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 + python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 - name: Run test low latency timeout-minutes: 10 diff --git a/tests/python/deepep/test_low_latency.py b/tests/python/deepep/test_low_latency.py index bd2869b83..6bebad3f5 100644 --- a/tests/python/deepep/test_low_latency.py +++ b/tests/python/deepep/test_low_latency.py @@ -75,17 +75,19 @@ def test( (num_local_experts,), dtype=torch.int, device="npu" ) 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, + 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 From 3f8e1f8002a3ea31ab7fd3e8b450a30bc92b36d2 Mon Sep 17 00:00:00 2001 From: zhuyy Date: Thu, 8 Jan 2026 21:03:15 +0800 Subject: [PATCH 5/5] delete little bs of intranode --- .github/workflows/pr-test-npu.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/pr-test-npu.yml b/.github/workflows/pr-test-npu.yml index c63300074..4ee7e3f2d 100644 --- a/.github/workflows/pr-test-npu.yml +++ b/.github/workflows/pr-test-npu.yml @@ -58,7 +58,6 @@ jobs: HCCL_BUFFSIZE: 2300 run: | python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py - python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 - name: Run test multi-round intranode timeout-minutes: 10 @@ -164,7 +163,6 @@ jobs: HCCL_BUFFSIZE: 2300 run: | python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py - python3 $GITHUB_WORKSPACE/tests/python/deepep/test_intranode.py --num-tokens=2 - name: Run test low latency timeout-minutes: 10