fix little batchsize and int8 quant on ci

.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

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):