# DeepEP CUDA Graph Compatibility: Limitations and Workarounds

> Explore DeepEP CUDA graph compatibility limitations. Discover workarounds for intranode NVLink operations, CPU sync, host tensor ops, and more in this technical guide.

- Repository: [DeepSeek/DeepEP](https://github.com/deepseek-ai/DeepEP)
- Tags: deep-dive
- Published: 2026-04-25

---

**DeepEP supports CUDA graph capture only for intranode NVLink operations when you disable CPU-side synchronization via `num_worst_tokens` and avoid host-side tensor operations like `record_stream`.**

DeepEP is DeepSeek's high-performance communication library for Mixture-of-Experts (MoE) models, providing optimized *dispatch* and *combine* kernels for GPU clusters. While CUDA graphs can eliminate CPU launch overhead and provide deterministic latency for MoE routing, **DeepEP CUDA graph compatibility** is constrained by strict requirements that prohibit host-side synchronization and incompatible stream operations within the captured graph.

## What Blocks CUDA Graph Capture in DeepEP?

CUDA graph capture requires that all operations be purely device-side, with no host CPU waits or cross-stream dependencies that trigger implicit synchronization. DeepEP's default implementations violate these constraints in several scenarios.

### Host-Side Synchronization in Normal Intranode Dispatch

Standard intranode dispatch operations in [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py) conclude with **CPU-side synchronization** via `torch.cuda.synchronize()` or host-side flag checks to verify token counts. According to the source comments at lines 354–355, this synchronization makes the operation incompatible with CUDA graph capture because the host wait would cause the graph to stall when replayed.

### Tensor Stream Recording Restrictions

Calling `tensor.record_stream()` on tensors living within the capture stream is forbidden during graph recording, as it implicitly creates a host-side dependency. The [`deep_ep/utils.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/utils.py) file addresses this limitation in lines 16–31, where the `EventOverlap` class documentation notes that stream recording via `record_stream` must be avoided to maintain **DeepEP CUDA graph compatibility**.

### Internode RDMA Coordination

Internode kernels always involve **RDMA/CPU coordination** for metadata exchange and host-side progress tracking. When `runtime.get_num_rdma_ranks() > 1`, the `Buffer.dispatch` method selects the internode path, which forces unavoidable host-side synchronization. This architectural requirement prevents internode kernels from being captured as-is.

### Low-Latency Kernel Defaults

Low-latency dispatch kernels do not synchronize the CPU-side token count by default, but the standard implementation performs a sync when the `recv_count` tensor is read on the host. As noted in [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py) at lines 94–95, this default behavior breaks CUDA graph compatibility unless explicitly configured otherwise.

## Workarounds for CUDA Graph Compatible Capture

To achieve successful graph capture, you must eliminate host-side dependencies and ensure all resources remain on the capture stream.

### Disable CPU Sync with num_worst_tokens

Set the `num_worst_tokens` argument to a value greater than zero (typically the batch size) when calling `dispatch()`. This forces the kernel to skip the CPU-side count synchronization, making the entire launch sequence graph-compatible. The `dispatch` method signature at line 327 of [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py) accepts this parameter specifically for CUDA graph scenarios.

### Avoid record_stream Using EventOverlap

Instead of calling `record_stream` directly, provide placeholder tensors via the `extra_tensors` argument of `EventOverlap`. These tensors are recorded on the current communication stream but are never accessed on the host, emulating the effect of `record_stream` without breaking capture. This pattern is implemented in [`deep_ep/utils.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/utils.py) at lines 19–31.

### Keep Operations on the Communication Stream

Pass `allocate_on_comm_stream=True` when creating `EventOverlap` instances or calling dispatch methods. This ensures all data buffers remain on the same stream used by the graph, preventing cross-stream synchronization barriers that would invalidate the capture.

### Restrict to Intranode Execution

CUDA graph capture is only supported for the **intranode** NVLink path. Internode kernels currently require host-side RDMA coordination that cannot be eliminated, making them unsuitable for graph capture until the library adds a fully asynchronous RDMA implementation.

## Code Examples

### Intranode Dispatch Inside a CUDA Graph

The following pattern shows how to capture a standard dispatch operation by combining `num_worst_tokens` with `EventOverlap`:

```python
import torch
import torch.distributed as dist
from deep_ep import Buffer, EventOverlap

# Prepare buffer for single-node NVLink execution

group = dist.new_group()
buffer = Buffer(group, num_nvl_bytes=0, num_rdma_bytes=0)

def graph_capture(x, topk_idx, topk_weights, num_experts):
    # Layout calculation must occur before graph capture (host ops)

    layout = buffer.get_dispatch_layout(
        topk_idx,
        num_experts,
        async_finish=True,
        allocate_on_comm_stream=True,
    )
    
    # Capture phase

    with torch.cuda.graph(torch.cuda.CUDAGraph()) as g:
        recv_x, recv_topk_idx, recv_topk_weights, _, handle, ev = buffer.dispatch(
            x,
            topk_idx=topk_idx,
            topk_weights=topk_weights,
            num_tokens_per_rank=layout.num_tokens_per_rank,
            num_tokens_per_rdma_rank=layout.num_tokens_per_rdma_rank,
            is_token_in_rank=layout.is_token_in_rank,
            num_tokens_per_expert=layout.num_tokens_per_expert,
            num_experts=num_experts,
            previous_event=None,
            async_finish=True,
            allocate_on_comm_stream=True,
            num_worst_tokens=x.size(0),  # Disables CPU sync for graph compatibility

        )
        # Supply extra tensors to safely record stream without host access

        event = EventOverlap(ev, extra_tensors=(x, topk_idx, recv_x))
    
    return g, (recv_x, recv_topk_idx, recv_topk_weights), handle, event

```

*Key implementation details:*
- `num_worst_tokens=x.size(0)` eliminates the host-side synchronization as referenced in [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py) line 354.
- `extra_tensors` satisfies stream recording requirements without breaking capture, per [`deep_ep/utils.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/utils.py) lines 19–31.

### Low-Latency Dispatch Capture

For low-latency kernels that avoid CPU count synchronization:

```python
def low_latency_graph(x, topk_idx, num_experts):
    with torch.cuda.graph(torch.cuda.CUDAGraph()) as g:
        recv_x, recv_cnt, handle, ev, hook = buffer.low_latency_dispatch(
            x,
            topk_idx,
            num_max_dispatch_tokens_per_rank=256,
            num_experts=num_experts,
            async_finish=True,
            return_recv_hook=False,
        )
        # No extra tensors needed when record_stream is avoided

        event = EventOverlap(ev)
    return g, recv_x, handle, event

```

This configuration leverages the fact that `low_latency_dispatch` does not synchronize the CPU received count with the GPU when `async_finish=True`, provided you never read `recv_count` on the host inside the graph.

## Summary

- **DeepEP CUDA graph compatibility** is restricted to intranode NVLink operations; internode RDMA kernels require host coordination that breaks capture.
- Set `num_worst_tokens > 0` in `Buffer.dispatch()` to disable CPU-side token count synchronization, as implemented in [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py) lines 327–356.
- Use `EventOverlap` with `extra_tensors` to simulate `record_stream` behavior without violating graph constraints, utilizing the pattern in [`deep_ep/utils.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/utils.py) lines 16–31.
- Always specify `allocate_on_comm_stream=True` to keep buffer allocations on the capture stream and prevent cross-stream dependencies.
- For low-latency dispatch, combine `async_finish=True` with host-avoidance of `recv_count` reads to maintain graph compatibility.

## Frequently Asked Questions

### Why do internode kernels break CUDA graph capture?

Internode kernels require RDMA metadata exchange and CPU-side progress coordination to manage data transfer between nodes. According to the DeepEP architecture, when `runtime.get_num_rdma_ranks() > 1`, the dispatch operation selects the internode path, which inherently involves host-side synchronization that CUDA graphs cannot record or replay.

### What is the purpose of extra_tensors in EventOverlap?

The `extra_tensors` parameter in `EventOverlap` provides a workaround for CUDA graphs' prohibition of `record_stream()` calls. By passing tensor references to this argument, as shown in [`deep_ep/utils.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/utils.py) lines 19–31, you ensure those tensors are associated with the current stream without triggering the host-side dependencies that `record_stream` normally creates, thereby maintaining graph capture validity.

### When should I use async_finish with DeepEP?

Set `async_finish=True` when you need CUDA graph compatibility or when chaining multiple communication operations to maximize overlap. This parameter returns a CUDA event instead of blocking the host, allowing subsequent operations to depend on completion via the event system rather than CPU synchronization. It is required for capturing both standard and low-latency dispatch kernels in graphs.

### Is low_latency_dispatch always graph-compatible?

No, `low_latency_dispatch` is only graph-compatible when you avoid reading the `recv_count` tensor on the host within the captured graph. While the kernel itself does not synchronize the CPU count by default, accessing this value on the CPU forces a host-side wait that breaks capture. Use `async_finish=True` and treat the operation as fire-and-forget to maintain compatibility.