How to Set Up DeepEP Auto-Tuning for Different Cluster Sizes

DeepEP requires re-tuning its static configuration map when deploying on clusters with different GPU counts or hardware topologies to achieve optimal dispatch and combine kernel performance.

DeepEP includes a pre-defined configuration table optimized for specific EP group sizes, but these defaults are calibrated on DeepSeek's internal cluster architecture. When running on different hardware—whether that means varying node counts, GPU models like A100 versus H100, or distinct NVLink topologies—you must benchmark the actual kernel performance on your target system and override the default config_map entries with the best-discovered parameters.

Understanding the Static Configuration Map

The Buffer class in deep_ep/buffer.py maintains a static config_map that defines dispatch and combine kernel launch configurations for common EP group sizes: 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 144, and 160. These configurations specify thread block layouts, SM utilization, and communication patterns.

While these defaults provide reasonable performance out-of-the-box, they cannot account for hardware-specific variables such as:

• NVLink versus PCIe bandwidth ratios on your specific nodes
• GPU architecture differences (SM counts vary between A100 40GB, A100 80GB, and H100)
• RDMA network topology for inter-node communication

For production deployments, treat these static entries as starting points rather than optimal settings.

Step-by-Step Auto-Tuning Workflow

Step 1: Deploy DeepEP on the Target Cluster

Install the DeepEP Python wheel or build from source, ensuring all processes can communicate via torch.distributed. Verify that torch.distributed.init_process_group(backend="nccl") initializes successfully across your target world_size.

Step 2: Run the Benchmark Test Suite

Execute the test suite on the exact number of EP ranks you plan to use in production. The repository provides three key test files that exercise different communication patterns:

• tests/test_intranode.py – Benchmarks intra-node dispatch/combine via NVLink
• tests/test_internode.py – Benchmarks inter-node dispatch/combine via RDMA
• tests/test_low_latency.py – Benchmarks low-latency kernels for inference scenarios

These tests iterate through every configuration entry in the internal config_map, printing elapsed times for each kernel variant. The benchmark utilities in deep_ep/utils.py provide the timing infrastructure used to measure latency.

Step 3: Collect the Best Configuration

Identify the configuration that achieved the lowest combined latency for dispatch and combine operations at your specific world_size. Save this configuration to a JSON file (e.g., dispatch_cfg_8.json for an 8-rank group) with the following structure:

{
  "num_sms": 24,
  "config": {
    "num_blocks": 128,
    "num_threads": 256
  }
}

Step 4: Inject the Auto-Tuned Configuration

You have two methods to force DeepEP to use your tuned configuration at runtime:

Method A: Explicit Configuration Passing

Pass the Config object directly to every dispatch and combine call:

buffer.dispatch(x, topk_idx=indices, topk_weights=weights, config=tuned_cfg)
buffer.combine(y, handle=recv_handle, config=tuned_cfg)

Method B: Monkey-Patch the Lookup Functions

Override the static methods Buffer.get_dispatch_config and Buffer.get_combine_config to return your tuned configuration automatically:

# Load the tuned config from file

with open(f"dispatch_cfg_{world_size}.json") as f:
    data = json.load(f)
    tuned_cfg = Buffer.Config(**data["config"])

# Replace the lookup functions globally

Buffer.get_dispatch_config = staticmethod(lambda _: tuned_cfg)
Buffer.get_combine_config = staticmethod(lambda _: tuned_cfg)

Method B ensures all subsequent buffer operations use the optimized settings without modifying individual call sites.

Complete Auto-Tuning Implementation

Below is a complete script that automates the benchmarking and configuration capture process. Deploy this via torch.distributed.launch with your target GPU count:

# auto_tune.py

import json
import torch
import torch.distributed as dist
from deep_ep import Buffer, EventOverlap
from deep_ep.utils import bench  # benchmark helper from tests

def _run_one_cfg(cfg, group):
    """Benchmark a single configuration."""
    hidden = 4096
    tokens = 256
    x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda")
    topk_idx = torch.randint(0, group.size() * 2, (tokens, 2), device="cuda", dtype=torch.int64)
    topk_weights = torch.randn_like(topk_idx, dtype=torch.float)

    buf = Buffer(group, 0, 0)
    
    # Benchmark dispatch

    dispatch_args = buf.get_dispatch_layout(topk_idx, num_experts=group.size() * 2)
    t_disp, _ = bench(
        lambda: buf.dispatch(
            x, topk_idx=topk_idx, topk_weights=topk_weights, config=cfg,
            **{k: v for k, v in zip(
                ["num_tokens_per_rank", "num_tokens_per_rdma_rank", 
                 "num_tokens_per_expert", "is_token_in_rank"], dispatch_args) if v is not None}
        )
    )
    
    # Benchmark combine

    recv_x, _, _ = buf.dispatch(x, topk_idx=topk_idx, topk_weights=topk_weights,
                                config=cfg, previous_event=None, async_finish=False)
    t_comb, _ = bench(lambda: buf.combine(recv_x, handle=recv_x.handle, config=cfg))
    
    return t_disp + t_comb

def auto_tune(group):
    """Find the fastest configuration for the current world size."""
    best_cfg = None
    best_time = float("inf")
    world = group.size()

    # Iterate through config_map entries for this world size

    for cfg in Buffer.get_dispatch_config(world).__class__.values():
        t = _run_one_cfg(cfg, group)
        if t < best_time:
            best_time, best_cfg = t, cfg
            
    return best_cfg

if __name__ == "__main__":
    dist.init_process_group(backend="nccl")
    group = dist.group.WORLD
    tuned_cfg = auto_tune(group)
    
    # Rank 0 saves the configuration

    if dist.get_rank() == 0:
        with open(f"dispatch_cfg_{group.size()}.json", "w") as fp:
            json.dump(
                {"num_sms": Buffer.num_sms, "config": tuned_cfg.__dict__},
                fp, indent=2
            )
    
    # Broadcast to all ranks and patch the lookup functions

    if dist.get_rank() == 0:
        with open(f"dispatch_cfg_{group.size()}.json") as fp:
            data = json.load(fp)
            tuned_cfg = Buffer.Config(**data["config"])
    else:
        tuned_cfg = None
        
    tuned_cfg = dist.broadcast_object_list([tuned_cfg], src=0)[0]
    Buffer.get_dispatch_config = staticmethod(lambda _: tuned_cfg)
    Buffer.get_combine_config = staticmethod(lambda _: tuned_cfg)

Run this script across your target cluster:

python -m torch.distributed.launch --nproc_per_node=8 auto_tune.py

Production Deployment Pattern

After generating the configuration files, load them in your production inference or training script:

import torch.distributed as dist
from deep_ep import Buffer
import json

dist.init_process_group(backend="nccl")
group = dist.group.WORLD

# Load the tuned configuration

with open(f"dispatch_cfg_{group.size()}.json") as f:
    cfg_dict = json.load(f)
tuned_cfg = Buffer.Config(**cfg_dict["config"])

# Apply globally

Buffer.get_dispatch_config = staticmethod(lambda _: tuned_cfg)
Buffer.get_combine_config = staticmethod(lambda _: tuned_cfg)

# Initialize buffer - will now use optimized settings

buffer = Buffer(group, num_nvl_bytes=..., num_rdma_bytes=...)

Advanced Tuning: SM Count and Low-Latency Mode

Tuning Streaming Multiprocessor (SM) Allocation

The static variable Buffer.num_sms controls how many streaming multiprocessors the high-throughput kernels utilize. Override this before creating any Buffer instances to optimize for your specific GPU model:

from deep_ep import Buffer
Buffer.set_num_sms(24)  # Must be an even number supported by your GPU

Benchmark different SM counts (typical values: 20, 24, 28, 32) to find the sweet spot between kernel occupancy and SM availability for other workloads.

Low-Latency Mode Considerations

When using low-latency mode (low_latency_mode=True), the optimal num_qps_per_rank (Queue Pairs per rank) depends on the number of local experts per GPU. Run tests/test_low_latency.py to tune these parameters for inference scenarios where minimal dispatch latency is critical.

Key Source Files Reference

The following files in the deepseek-ai/DeepEP repository contain the implementation details referenced above:

• deep_ep/buffer.py – Contains the Buffer class, static config_map, get_dispatch_config, get_combine_config, and set_num_sms methods
• tests/test_intranode.py – Intra-node benchmark suite for NVLink configurations
• tests/test_internode.py – Inter-node benchmark suite for RDMA configurations
• tests/test_low_latency.py – Low-latency kernel benchmarks for inference optimization
• deep_ep/utils.py – Timing utilities including the bench function

Summary

• DeepEP's default config_map in deep_ep/buffer.py covers common EP group sizes (2–160 ranks) but requires re-tuning for different hardware topologies.
• Run tests/test_intranode.py, tests/test_internode.py, and tests/test_low_latency.py on your target cluster to benchmark every configuration entry.
• Save the best-performing Config object for your specific world_size and inject it either by passing it explicitly to dispatch/combine calls or by overriding Buffer.get_dispatch_config and Buffer.get_combine_config.
• Tune Buffer.num_sms via Buffer.set_num_sms() to optimize SM utilization for your specific GPU model (A100, H100, etc.).
• For inference workloads using low-latency mode, validate num_qps_per_rank settings against your specific expert parallelism configuration.

Frequently Asked Questions

What cluster sizes require custom auto-tuning?

Any cluster configuration that differs from DeepSeek's internal topology—including different GPU models (A100 vs H100), varying numbers of GPUs per node, or distinct NVLink/RDMA network fabrics—requires re-tuning. Even if your EP group size (e.g., 8 ranks) exists in the static config_map, the optimal thread block layout and SM count likely differ on your hardware.

Can I skip auto-tuning and use the default configurations?

You can use the static defaults for initial testing, but production deployments should always run the full benchmark suite. The default configurations prioritize generality over peak performance and may utilize suboptimal SM counts or thread layouts for your specific NVLink topology or RDMA latency characteristics.

How do I tune DeepEP for low-latency inference mode?

First, ensure you are running tests/test_low_latency.py rather than the standard throughput tests. Low-latency mode uses different kernel implementations that require specific num_qps_per_rank values matching your local expert count. The auto-tuning process remains similar: benchmark all candidate configurations, save the fastest, and inject via the config parameter when creating low-latency buffers.

Why does the SM count setting affect performance?

The Buffer.num_sms variable controls how many streaming multiprocessors are dedicated to the dispatch and combine kernels. Setting this too high can starve other GPU operations, while setting it too low leaves compute resources idle. The optimal value varies by GPU architecture—H100 GPUs have different SM counts and memory bandwidth characteristics compared to A100s—making per-hardware tuning essential for maximizing throughput.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}