[GH-ISSUE #2087] [Question]: Does Mooncake EP currently support CUDA Graph at all? #571

New issue

Closed

opened 2026-10-29 21:16:38 +00:00 by bsakaguchi · 3 comments

bsakaguchi commented

2026-10-29 21:16:38 +00:00

Owner

Originally created by @Socratesa on GitHub (May 12, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/2087

Originally assigned to: @UNIDY2002 on GitHub.

Bug Report

I've been studying the Mooncake EP source code and plan to integrate it into a custom MoE inference engine. I have a question about CUDA Graph compatibility.

Looking at mooncake_ep_buffer.cpp, each dispatch() / combine() call allocates temporary tensors like packed_recv_x, packed_recv_src_info, packed_recv_count via torch::empty() (L248-L259), and in async mode creates CUDA Events for cross-stream synchronization (L306 EventHandle(launch_stream)).

I also noticed that the Python-level EventOverlap class has already introduced extra_tensors as a workaround, with comments noting:

stream recording will be incompatible with CUDA graph we use extra tensors to simulate PyTorch tensor record_stream, may be useful with CUDA graph

A few questions:

Does Mooncake EP currently support CUDA Graph at all?
Is there a plan to add CUDA Graph support?

Before submitting...

Ensure you searched for relevant issues and read the [documentation]

Originally created by @Socratesa on GitHub (May 12, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/2087 Originally assigned to: @UNIDY2002 on GitHub. ### Bug Report I've been studying the Mooncake EP source code and plan to integrate it into a custom MoE inference engine. I have a question about CUDA Graph compatibility. Looking at mooncake_ep_buffer.cpp, each dispatch() / combine() call allocates temporary tensors like packed_recv_x, packed_recv_src_info, packed_recv_count via torch::empty() (L248-L259), and in async mode creates CUDA Events for cross-stream synchronization (L306 EventHandle(launch_stream)). I also noticed that the Python-level EventOverlap class has already introduced extra_tensors as a workaround, with comments noting: stream recording will be incompatible with CUDA graph we use extra tensors to simulate PyTorch tensor record_stream, may be useful with CUDA graph A few questions: 1. Does Mooncake EP currently support CUDA Graph at all? 2. Is there a plan to add CUDA Graph support? ### Before submitting... - [x] Ensure you searched for relevant issues and read the [documentation]

bsakaguchi

2026-10-29 21:16:38 +00:00

closed this issue
added the
stale

auto-closed

bug
labels

bsakaguchi commented

2026-10-29 21:16:39 +00:00

Author

Owner

@UNIDY2002 commented on GitHub (May 13, 2026):

Hi @Socratesa, thanks for the detailed analysis. Mooncake EP does work with CUDA graph — this is the primary production path in SGLang and has been tested. Your concerns are valid but depend on which flags are used.

Regarding each point:

torch::empty() allocations: PyTorch patches the CUDA caching allocator during torch.cuda.graph() capture to use the graph's private memory pool. Since decode has a fixed token count, the allocation sizes are consistent across replays.
CUDA Events for cross-stream synchronization: This only happens when return_recv_hook=False, where dispatch()/combine() use a separate comm_stream and call stream_wait() (which involves cudaEventRecord + cudaStreamWaitEvent). When return_recv_hook=True, launch_stream = compute_stream (the same stream being captured), and no cross-stream synchronization occurs.
extra_tensors: Your understanding is correct — it serves the same purpose as record_stream() and happens to be compatible with CUDA graph since it doesn't rely on the stream recording mechanism.

The key is using the right flags. For CUDA graph capture, you should set:

return_recv_hook=True — keeps all kernel launches on the compute stream (the one being captured)
async_finish=False — required when return_recv_hook=True (enforced by an assert in the C++ code)

@UNIDY2002 commented on GitHub (May 13, 2026): Hi @Socratesa, thanks for the detailed analysis. Mooncake EP does work with CUDA graph — this is the primary production path in SGLang and has been tested. Your concerns are valid but depend on which flags are used. Regarding each point: 1. **`torch::empty()` allocations**: PyTorch patches the CUDA caching allocator during `torch.cuda.graph()` capture to use the graph's private memory pool. Since decode has a fixed token count, the allocation sizes are consistent across replays. 2. **CUDA Events for cross-stream synchronization**: This only happens when `return_recv_hook=False`, where `dispatch()`/`combine()` use a separate `comm_stream` and call `stream_wait()` (which involves `cudaEventRecord` + `cudaStreamWaitEvent`). When `return_recv_hook=True`, `launch_stream = compute_stream` (the same stream being captured), and no cross-stream synchronization occurs. 3. **`extra_tensors`**: Your understanding is correct — it serves the same purpose as `record_stream()` and happens to be compatible with CUDA graph since it doesn't rely on the stream recording mechanism. The key is using the right flags. For CUDA graph capture, you should set: - **`return_recv_hook=True`** — keeps all kernel launches on the compute stream (the one being captured) - **`async_finish=False`** — required when `return_recv_hook=True` (enforced by an assert in the C++ code)

bsakaguchi commented

2026-10-29 21:16:39 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Aug 12, 2026):

This issue has had no activity for 90 days and will be closed in 7 days if there is no further activity. Please comment or react if it should stay open.

@github-actions[bot] commented on GitHub (Aug 12, 2026): This issue has had no activity for 90 days and will be closed in 7 days if there is no further activity. Please comment or react if it should stay open.

bsakaguchi commented

2026-10-29 21:16:39 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Aug 19, 2026):

Closing due to 3 months of inactivity. If this is still relevant, please comment and we can reopen.

@github-actions[bot] commented on GitHub (Aug 19, 2026): Closing due to 3 months of inactivity. If this is still relevant, please comment and we can reopen.

bsakaguchi referenced this issue

2026-10-29 21:29:39 +00:00

[PR #633] [CLOSED] [do not merge]: add g++ flags for (#571) to allow proper arm build #1704