diff --git a/benchmark/kernels/deepep/tuning_deepep.py b/benchmark/kernels/deepep/tuning_deepep.py index 191819d2c..bab37b848 100644 --- a/benchmark/kernels/deepep/tuning_deepep.py +++ b/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,): diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 8539639d5..f6e35e401 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/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