fix(deepep): tuning and enhance deepep on single node

feat(deepep): enhance single node deepep

benchmark/kernels/deepep/tuning_deepep.py

@@ -46,6 +46,10 @@ def test_main(
         args.num_topk,
         (args.num_experts // num_ranks) * num_ranks,
     )
+    if args.num_experts is not None:
+        num_experts = args.num_experts
+    else:
+        num_experts = (256 // num_ranks) * num_ranks
     assert num_experts % num_ranks == 0 and num_local_ranks == 8
     if local_rank == 0:
         print(
@@ -125,7 +129,8 @@ def test_main(
         _,
     ) = buffer.get_dispatch_layout(topk_idx, num_experts)
     assert torch.allclose(ref_num_tokens_per_rank, num_tokens_per_rank)
-    assert torch.allclose(ref_num_tokens_per_rdma_rank, num_tokens_per_rdma_rank)
+    if ref_num_tokens_per_rdma_rank is not None:
+        assert torch.allclose(ref_num_tokens_per_rdma_rank, num_tokens_per_rdma_rank)
     assert torch.allclose(ref_num_tokens_per_expert, num_tokens_per_expert)
     assert torch.allclose(ref_is_token_in_rank, is_token_in_rank)
     t = bench(lambda: buffer.get_dispatch_layout(topk_idx, num_experts))[0]
@@ -136,7 +141,16 @@ def test_main(
     time.sleep(1)
 
     # Config
-    rdma_buffer_size, nvl_buffer_size = 128, (720 if num_ranks in (144, 160) else 512)
+    rdma_buffer_size = 128
+    # Calculate nvl_buffer_size to be divisible by num_nodes (num_rdma_ranks)
+    base_nvl_buffer_size = 720 if num_ranks in (144, 160) else 512
+    nvl_buffer_size = ((base_nvl_buffer_size + num_nodes - 1) // num_nodes) * num_nodes
+    if local_rank == 0:
+        print(
+            f"[config] num_nodes={num_nodes}, base_nvl_buffer_size={base_nvl_buffer_size}, "
+            f"nvl_buffer_size={nvl_buffer_size}, divisible_check={nvl_buffer_size % num_nodes}",
+            flush=True,
+        )
     config = deep_ep.Config(num_sms, 8, nvl_buffer_size, 16, rdma_buffer_size)
 
     # Test dispatch
@@ -197,7 +211,10 @@ def test_main(
                     )
 
                     # Checks
-                    recv_gbl_rank_prefix_sum = handle[-4]
+                    if num_nodes == 1:
+                        recv_gbl_rank_prefix_sum = handle[0][:, rank]
+                    else:
+                        recv_gbl_rank_prefix_sum = handle[-4]
                     assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
                         0
                     ), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
@@ -312,8 +329,9 @@ def test_main(
             if isinstance(current_x, tuple)
             else dispatch_bf16_nvl_recv_bytes
         )
+        rdma_chunk_range = range(4, 33, 4) if num_nodes > 1 else (32,)
         for nvl_chunk_size in range(4, 33, 4):
-            for rdma_chunk_size in range(4, 33, 4):
+            for rdma_chunk_size in rdma_chunk_range:
                 config_kwargs = {
                     "num_sms": num_sms,
                     "num_max_nvl_chunked_send_tokens": nvl_chunk_size,
@@ -381,8 +399,9 @@ def test_main(
 
     # Tune combine performance
     best_time, best_results = 1e10, None
+    rdma_chunk_range = range(12 if num_nodes == 2 else 8, 33, 4) if num_nodes > 1 else (32,)
     for nvl_chunk_size in range(1, 8, 1):
-        for rdma_chunk_size in range(12 if num_nodes == 2 else 8, 33, 4):
+        for rdma_chunk_size in rdma_chunk_range:
             config_kwargs = {
                 "num_sms": num_sms,
                 "num_max_nvl_chunked_send_tokens": nvl_chunk_size,
@@ -431,16 +450,21 @@ def test_loop(local_rank: int, num_local_ranks: int, args):
     rank, num_ranks, group = init_dist(local_rank, num_local_ranks, args)
 
     num_sms = args.num_sms
-    num_qps_per_rank = num_sms // 2
+    if num_nodes == 1:
+        num_rdma_bytes = 0
+        num_qps_per_rank = 1
+    else:
+        num_rdma_bytes = int(1e9)
+        num_qps_per_rank = num_sms // 2
 
     buffer = deep_ep.Buffer(
         group,
         int(1e9),
-        int(1e9),
+        num_rdma_bytes,
         low_latency_mode=False,
         num_qps_per_rank=num_qps_per_rank,
     )
-    assert num_local_ranks == 8 and num_ranks > 8
+    assert num_local_ranks == 8 and num_ranks >= 8
     torch.manual_seed(rank)
 
     for i in (num_sms,):

python/sglang/srt/layers/moe/token_dispatcher/deepep.py

@@ -218,14 +218,19 @@ class DeepEPBuffer:
                     f"Consider using --deepep-config to change the behavior."
                 )
 
+        is_single_node = group.size() <= torch.cuda.device_count()
+        allow_mnnvl = True
+        if is_single_node and not deepep_mode.enable_low_latency():
+            num_rdma_bytes = 0
+            allow_mnnvl = False
+
         cls._buffer = Buffer(
             group,
             num_nvl_bytes,
             num_rdma_bytes,
             low_latency_mode=deepep_mode.enable_low_latency(),
             num_qps_per_rank=num_qps_per_rank,
-            # TODO can be false when unneeded
-            allow_mnnvl=True,
+            allow_mnnvl=allow_mnnvl,
         )
         return cls._buffer