DeepEP Low‑Latency Dispatch and Combine Kernels for Inference Decoding

DeepEP provides specialized RDMA‑based dispatch and combine kernels that enable sub‑millisecond token routing for MoE inference decoding, implemented in the deep_ep.Buffer class and backed by the C++ runtime.

The deepseek-ai/DeepEP repository delivers high‑performance communication primitives for Mixture‑of‑Experts (MoE) models. Its low‑latency dispatch and combine kernels are engineered specifically for the tight response‑time requirements of inference decoding, where tokens must be routed to experts and recombined with minimal overhead.

Core Architecture and Transport

DeepEP’s low‑latency kernels rely on pure RDMA All‑to‑All communication using IBGDA (in‑band gather‑direct‑access). This design eliminates CPU intervention and PCIe bottlenecks, allowing GPUs to write directly to remote memory. All participating ranks must be reachable via RDMA regardless of node topology. The Python wrapper class deep_ep.Buffer manages the lifecycle of these kernels, forwarding heavy computation to the underlying deep_ep_cpp runtime.

Both kernels employ a minimal buffering strategy—only two shared buffers are reused across calls—to reduce memory pressure and allocation latency.

Low‑Latency Dispatch Kernel

The dispatch kernel sends each token to the expert(s) selected by a top‑k routing mask, producing per‑expert tensors ready for the next compute step.

RDMA Requirements and Validation

Before initiating communication, the kernel validates that the RDMA queue‑pair depth can accommodate the token volume. In deep_ep/buffer.py, the assertion ensures sufficient headroom for the maximum dispatch size:

assert self.nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2

This check prevents queue‑pair exhaustion during high‑throughput inference batches.

API Signature and Return Values

The low_latency_dispatch method accepts the input tensor, routing indices, and configuration flags, returning a tuple containing the packed receive buffer, counts, and synchronization primitives:

packed_recv_x, packed_recv_count, handle, event, hook = buf.low_latency_dispatch(
    x, topk_idx,
    num_max_dispatch_tokens_per_rank=256,
    num_experts=64,
    use_fp8=True,
    round_scale=False,
    use_ue8m0=False,
    async_finish=False,
    return_recv_hook=False,
    cumulative_local_expert_recv_stats=None,
    dispatch_wait_recv_cost_stats=None
)

The handle bundles metadata required for the subsequent combine operation: source information, layout ranges, and tensor dimensions.

FP8 and Quantization Support

The dispatch kernel supports FP8 (torch.float8_e4m3fn) activation transfer via the use_fp8 flag. When enabled, the system returns a tuple (packed_recv_x, packed_recv_x_scales) where scales can be stored in UE8M0 format (use_ue8m0=True). This reduces RDMA bandwidth by 50% compared to BF16, critical for latency‑serving scenarios.

Optional statistics collection tracks per‑expert receive counters (cumulative_local_expert_recv_stats) and detailed wait‑time histograms (dispatch_wait_recv_cost_stats) for performance analysis.

Low‑Latency Combine Kernel

The combine kernel performs the inverse operation: it reduces (sums) the per‑expert results back to the original token order, applying the top‑k weights to produce the final hidden states.

Reduction Logic and Buffer Reuse

Mirroring the dispatch architecture, combine reuses the dual‑buffer strategy and validates queue‑pair depth against the same num_max_dispatch_tokens_per_rank threshold. The C++ runtime method runtime.low_latency_combine handles the actual reduction, accessed via:

src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
combined_x, event, hook = self.runtime.low_latency_combine(
    x, topk_idx, topk_weights,
    src_info, layout_range,
    combine_wait_recv_cost_stats,
    num_max_dispatch_tokens_per_rank, num_experts,
    use_logfmt=False,
    zero_copy=False,
    async_finish=False,
    return_recv_hook=False,
    out=None
)

Zero‑Copy Optimization

For advanced use cases, the combine kernel supports a zero‑copy mode that writes directly into a pre‑registered RDMA buffer. When zero_copy=True, the caller must first acquire the destination buffer via get_next_low_latency_combine_buffer, eliminating an intermediate copy and shaving microseconds off the critical path.

Practical Implementation Guide

The following end‑to‑end example demonstrates a complete forward pass using the low‑latency APIs. This pattern follows the integration tests in tests/test_low_latency.py.

import torch
import deep_ep
from deep_ep.utils import topk_idx_t

# Initialize the EP buffer in low‑latency mode

group = torch.distributed.new_group()
buf = deep_ep.Buffer(
    group=group,
    num_nvl_bytes=0,
    num_rdma_bytes=1 << 26,
    low_latency_mode=True,
    num_qps_per_rank=24,
    allow_nvlink_for_low_latency_mode=False
)

# Prepare input tokens and routing

num_tokens, hidden, num_experts, topk = 1024, 4096, 64, 2
x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device='cuda')
scores = torch.randn(num_tokens, num_experts, device='cuda')
topk_idx = torch.topk(scores, topk, dim=-1).indices.to(topk_idx_t)
topk_weights = torch.rand(num_tokens, topk, device='cuda')

# Dispatch: send tokens to expert‑local buffers

packed_recv_x, packed_recv_count, handle, event, hook = buf.low_latency_dispatch(
    x, topk_idx,
    num_max_dispatch_tokens_per_rank=256,
    num_experts=num_experts,
    use_fp8=True,
    async_finish=False,
    return_recv_hook=False
)
event.current_stream_wait()

# Unpack FP8 if needed and simulate expert computation

if isinstance(packed_recv_x, tuple):
    recv_x = deep_ep.utils.per_token_cast_back(packed_recv_x[0], packed_recv_x[1])
else:
    recv_x = packed_recv_x

expert_out = recv_x @ torch.randn(hidden, hidden, dtype=torch.bfloat16, device='cuda')

# Combine: reduce results back to token order

combined_x, combine_event, combine_hook = buf.low_latency_combine(
    expert_out, topk_idx, topk_weights,
    handle,
    zero_copy=False,
    async_finish=False,
    return_recv_hook=False
)
combine_event.current_stream_wait()

print("Combined output shape:", combined_x.shape)  # [num_tokens, hidden]

Performance Optimizations and Failure Resilience

DeepEP’s low‑latency kernels expose several mechanisms to minimize synchronization overhead and handle partial failures:

• Asynchronous Execution: Setting async_finish=True enables CUDA‑graph‑compatible non‑blocking execution, allowing the host to post successive operations while RDMA transfers proceed in the background.
• Early Return Hooks: When return_recv_hook=True, the dispatch kernel issues RDMA reads immediately but returns a hook object that the caller can wait on later, hiding latency behind computation.
• Rank Masking: A rank can mask itself (or be masked) using low_latency_update_mask_buffer. The kernel skips communication with masked ranks and leaves buffers clean. The current mask state is queried via low_latency_query_mask_buffer, enabling graceful degradation when individual nodes become overloaded or unreachable.

Summary

• DeepEP implements inference‑optimized dispatch and combine kernels in the deep_ep.Buffer class, delegating to the deep_ep_cpp runtime.
• Both kernels use IBGDA‑based RDMA All‑to‑All for sub‑millisecond latency, requiring nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2.
• Dispatch sends tokens to experts, supporting FP8 with optional UE8M0 scaling, while combine reduces results using top‑k weights with optional zero‑copy output.
• Dual shared buffers minimize allocation overhead, and optional hooks (return_recv_hook) enable asynchronous completion patterns.
• Rank masking APIs allow dynamic exclusion of failed or slow nodes without restarting the communication group.

Frequently Asked Questions

How does DeepEP achieve low latency compared to standard NCCL collectives?

DeepEP bypasses the CPU entirely by using IBGDA (in‑band gather‑direct‑access), which allows GPUs to initiate RDMA writes directly to remote GPU memory. Standard NCCL all‑to‑all often requires CPU coordination and additional PCIe copies, whereas DeepEP’s kernels in deep_ep/buffer.py validate queue‑pair depth once and reuse two shared buffers across calls, eliminating allocation latency and kernel‑launch overhead.

What is the difference between the dispatch and combine handles?

The dispatch handle is a tuple (packed_recv_src_info, packed_recv_layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts) that encodes the routing metadata necessary for the combine phase. This handle must be passed unchanged to low_latency_combine so the kernel knows how to map per‑expert results back to the original token indices and apply the top‑k weights correctly.

When should I use FP8 mode in the dispatch kernel?

Enable use_fp8=True when bandwidth is the bottleneck rather than compute, typically in large‑scale inference deployments where tensor sizes approach the RDMA bandwidth limit. The kernel supports torch.float8_e4m3fn and optional UE8M0 scaling (use_ue8m0=True), reducing inter‑node traffic by 50%. After expert computation, use deep_ep.utils.per_token_cast_back to recover BF16 precision before the combine step.

How does the zero‑copy mode in combine improve performance?

Zero‑copy mode (zero_copy=True) allows the combine kernel to write results directly into a pre‑registered RDMA buffer obtained via get_next_low_latency_combine_buffer. This eliminates an extra device‑to‑device copy that would otherwise occur when transferring data from the internal packed buffer to the final output tensor, saving approximately 5–10 microseconds in latency‑critical decoding loops.

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