[EP] decouple native EP from torch C++ APIs#2883
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the Mooncake EP backend to remove its C++ dependency on PyTorch, transitioning tensor checks and memory management to the Python layer and exposing C++ interfaces via raw pointers and primitive types. The build system is also updated to compile the backend as a standard Pybind11 module. The review feedback highlights several critical performance and safety concerns, including synchronous cudaMalloc/cudaFree allocations on the hot path of dispatch, potential race conditions from manual reference counting in EventHandle, driver overhead from dynamic CUDA event creation in stream_wait, and potential multi-GPU issues from hardcoding device="cuda" in the Python fallback buffer allocation.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![qoderai[bot]](https://avatars.githubusercontent.com/in/1396158?s=60&v=4){kind=small}
There was a problem hiding this comment.
👋 Review Summary
This refactor effectively decouples the EP native layer from Torch C++ and moves validation into the Python wrappers, while integrating EP into the main CMake graph. The overall structure and tests look solid, and the changes align with the goal of making EP a first-class, Torch-agnostic backend.
🛡️ Key Risks & Issues
- Raw-pointer EP Buffer dispatch bug (mooncake_ep_buffer.cpp:214): In the new pointer-based dispatch path,
packed_recv_x_scalesis initialized from a localpacked_recv_x_scales_ptrvariable that is never assigned from the incoming uint64_t argument. This means the FP8 branch uses an undefined pointer even when Python passes a valid recv_x_scales tensor. Fixing the shadowed variable (wiring the function parameter directly into the reinterpret_cast) should restore correct behavior. - Elastic deterministic path allocations (mooncake_ep_elastic_buffer.cpp:290+): The deterministic dispatch path now calls cudaMalloc/cudaFree on every dispatch to allocate
deterministic_rank_count_buffer, which can add noticeable latency and allocator pressure under high QPS. Pre-allocating this buffer per ElasticBuffer instance (and reusing it via the existing keepalive mechanism) would reduce hot-path overhead while preserving determinism. - EventHandle API semantics and external callers (mooncake_ep_event.h:9+): EventHandle has changed from a Torch Event wrapper that records on the current stream and offers a parameterless current_stream_wait() to a raw CUDA event requiring explicit stream pointers. Internal callers (via EventOverlap and _native_current_stream_ptr) look correct, but any external code that previously used EventHandle() and current_stream_wait() still expects the old semantics. Clarifying or enforcing the new contract (e.g., asserting on stream_ptr==0, or preserving default-current-stream behavior) would avoid silent mis-synchronization.
- EP/PG active-ranks dependency (mooncake_ep_buffer.py:328+): Buffer._active_ranks_tensor now imports get_active_ranks from mooncake.pg. That’s a reasonable consolidation, but it couples EP to PG’s notion of group membership. If PG returns a mask whose length doesn’t match EP’s group size (or PG is misconfigured/absent), EP can silently fall back to all-ones or use a mismatched mask. A simple shape check before applying the PG mask would make misconfigurations easier to catch.
- Elastic Python-side CPU-sync and handle reuse (mooncake_elastic_buffer.py:480+): The elastic dispatch wrapper now handles prefix-sum interpretation and per-expert counts on the host when do_cpu_sync is enabled, and cached handles carry a boolean native_handle with Python-owned tensors for metadata. This design is coherent, but leaving do_cpu_sync on in production will add host copies and processing, and reusing a handle after its backing tensors change could lead to subtle inconsistencies. Adding a bit more validation around handle reuse (e.g., shape checks) and treating do_cpu_sync as a debug-only mode would reduce that risk.
🧪 Verification Advice
- Re-run the existing EP and elastic tests you listed, especially:
python3 mooncake-wheel/tests/test_mooncake_ep.pypython3 -m torch.distributed.run --nnodes=1 --nproc_per_node=2 mooncake-ep/tests/test_elastic_buffer.py ...
- After fixing the packed_recv_x_scales pointer wiring, add a focused FP8 EP test case that:
- Uses use_fp8=True, verifies recv_x_scales is populated correctly, and compares results against a known-good BF16 path.
- For the elastic deterministic path, consider a micro-benchmark (e.g., mooncake-ep/benchmarks/elastic_buffer_perf.py) configured with deterministic=True to confirm that moving allocations out of the hot path does not change correctness and improves tail latency.
💡 Thoughts & Suggestions
- The overall architectural direction—Python-managed tensors and raw-pointer native kernels—is sensible and aligns with removing Torch dependencies from the extension. To keep it maintainable, it will help to make the Python invariants (shape, dtype, group-size expectations) very explicit and assertive so future callers can’t accidentally bypass them.
- Given the new EP/PG coupling, adding a short note in the EP docs or code about how active_ranks is sourced from PG (and what assumptions are made about group alignment) would help future contributors reason about failures in distributed setups.
- If you plan to extend this approach to other backends (MUSA/MACA), you might want a small shared utility for mapping between torch streams and native cuda/musa streams, so EventHandle semantics remain consistent across devices.
🤖 Generated by Qoder • View workflow run
There was a problem hiding this comment.
👋 Review Summary
This refactor does a solid job of decoupling the EP native layer from torch C++ APIs, with clear pointer-based interfaces and Python wrappers handling shape, stream, and routing logic. The build changes and Python-backed tests you describe line up well with the new _ep module and the elastic buffer restructuring.
🛡️ Key Risks & Issues
- EP event bridging: the new
_wait_native_event_on_current_streamhelper expects events to implementcurrent_stream_wait(stream_ptr), and the nativeEventHandlenow matches that signature. However, the fallback_DummyEventstill definescurrent_stream_wait()with no arguments and is used as a drop-in replacement in fallback EP paths. When wrapped inEventOverlap, this will cause aTypeError(or mis-synchronization) on the dummy-event path rather than providing the intended compatibility, which can break EP under fallback/test configurations.
🧪 Verification Advice
- Add a focused unit test that constructs an
EventOverlaparound_DummyEventand exercisescurrent_stream_wait()and the context manager on a real CUDA stream, to confirm that the dummy-event path cleanly synchronizes without exceptions. - For the elastic buffer CPU-sync logic, consider adding small, property-style tests that validate
num_recv_tokens_per_expert_listand recv buffer slicing across differentexpert_alignment,num_experts, andnum_max_tokens_per_ranksettings, to catch any subtle off-by-one issues in the Python reimplementation.
💡 Thoughts & Suggestions
- For
_DummyEvent, either makecurrent_stream_waitaccept an optional stream pointer (and ignore it) or special-case it inside_wait_native_event_on_current_streamso fallback always usessynchronize(). That keeps the dummy-event path robust while preserving the more precise native stream behavior. More broadly, the pointer-based native APIs and explicit assertions in the Python wrappers look good and should help keep EP behavior predictable as you evolve the CUDA and hybrid topologies.
🤖 Generated by Qoder • View workflow run
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
👋 Review Summary
This is a substantial and well-scoped refactor that cleanly decouples the Mooncake EP/elastic native layer from Torch C++ APIs, moves the EP build into the main CMake graph, and modernizes the Python wrappers to own tensor allocation and stream management. The overall direction looks solid and the changes are thoughtfully defended with assertions and end-to-end tests.
🛡️ Key Risks & Issues
- Zero-copy combine semantics: the
Buffer.get_next_combine_bufferAPI no longer wires through to a native EP combine buffer, and the newBuffer.combineimplementation always uses a freshcombined_xtensor even whenzero_copy=True. This preserves correctness but effectively disables the original zero-copy fast path and may surprise callers that rely on the existing contract; I left an inline comment with details and suggestions. - Elastic handle / prefix-sum coupling: the elastic path now reconstructs
EPHandleentirely in Python from prefix-sum metadata. This is a reasonable tradeoff for Torch-decoupling, but it tightly couples Python logic to native layout assumptions (alignment, sentinel elements, hybrid routing). Future changes to the C++ kernels or metadata layout will need coordinated updates to the Python reconstruction to avoid subtle expert-count or truncation bugs. - Event / stream behavior: the new CUDA-only
EventHandleandEventOverlapintegration looks correct, but behavior when the current stream pointer is 0 or when callers use non-default streams isn’t explicitly exercised. Given how central these events are to overlapping compute and communication, it’s worth keeping an eye on any future changes here.
🧪 Verification Advice
- Rerun the existing EP and elastic tests under a variety of configurations, including:
- different
MOONCAKE_EP_NUM_LOCAL_RANKS/ world sizes, and bothroute=localandroute=crosspaths; deterministic=Trueon ElasticBuffer, to exercise the deterministic rank-count prologue and buffer reuse; and- empty or near-empty batches (e.g.,
num_tokens=0and smallnum_topk) to validate the pointer/null handling on the new EP dispatch/combine APIs.
- different
- If zero-copy combine performance is important for your workloads, consider a focused microbenchmark comparing the old behavior vs the new one and decide whether to reintroduce a true native buffer-based path or to update docs/tests to reflect the simpler copy-based semantics.
💡 Thoughts & Suggestions
- Overall this is a nice clean-up that should make EP/elastic more portable and easier to maintain without Torch C++ internals. To reduce future maintenance risk, it may help to (a) centralize the elastic metadata contract in a small shared description (e.g., comments or a docstring) referenced by both C++ and Python, and (b) add a dedicated test or microbenchmark for the intended fast paths (FP8 dispatch, cached elastic handles, and any future zero-copy combine variant) so regressions in those performance-sensitive areas are caught early.
🤖 Generated by Qoder • View workflow run
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
There was a problem hiding this comment.
👋 Review Summary
This refactor does a solid job of decoupling the native EP layer from Torch C++ while keeping the Python APIs coherent and well-tested. The main dispatch/combine flows for both legacy EP and Elastic EP look consistent and are backed by good test coverage, and the new native legacy benchmark is a useful addition.
🛡️ Key Risks & Issues
- Raw-pointer native EP APIs (
MooncakeEpBuffer::dispatch/combine) now rely almost entirely on Python wrappers for invariants. That’s fine for current usage, but it makes the C++ entrypoints easy to misuse from any future internal caller; adding a few pointer non-null/size sanity checks would make them more robust without reintroducing Torch dependencies. - The
EventHandlePython binding now requires explicit stream pointers for construction andcurrent_stream_wait, which can break external code that previously used zero-argument calls and makes it easier to accidentally pass the wrong CUDA stream handle. If this class is meant to stay public, a safer overload that uses the current stream would help. - Elastic dispatch has gained significant Python-side metadata logic (prefix sums, routing metadata, cached-handle reuse). The non-expanded BF16 path is well covered by
test_elastic_buffer.py, but FP8 dispatch,do_expand=True, and deterministic + cached-handle paths are not obviously exercised. Bugs there would mostly show up as subtle routing or numerical issues rather than hard failures. Buffer.get_next_combine_buffernow always returns a Python-allocated staging tensor and no longer exposes the native EP combine buffer. That’s consistent with the new pointer-based combine interface but subtly changes the zero-copy story at the Python layer and could mislead future users if old semantics are assumed.
🧪 Verification Advice
- Continue running the EP-focused test suite as in the PR description:
test_mooncake_ep.py,test_elastic_buffer.py, andtest_ep_grid.pyon representative hardware and interconnects. - Add targeted tests for: FP8 elastic dispatch (including scales),
do_expand=Trueelastic flows, and deterministic + cached-handle reuse, ideally validating both metadata (psum_*,recv_src_metadata) and payload correctness. - Consider a small unit test around
EventOverlapandEventHandleto exercise both native and fallback paths, ensuring stream-wait semantics remain correct if the event API changes again.
💡 Thoughts & Suggestions
- Overall the direction (Torch-free native layer, Python owning tensor semantics) makes sense and aligns with long-term decoupling goals. To keep it maintainable, I’d recommend documenting the intended usage patterns for the raw-pointer C++ APIs and clarifying in docs that
get_next_combine_buffer()is now a Python staging buffer rather than a transport-managed view. - If you expect external consumers to touch
EventHandleor Elastic internals, a short section in the EP docs describing the new stream-pointer-based API and the constraints on cached handles (non-expanded only, deterministic behavior) would reduce surprises.
🤖 Generated by Qoder • View workflow run
| combine(const torch::Tensor& x, const torch::Tensor& topk_idx, | ||
| const torch::Tensor& topk_weights, const torch::Tensor& src_info, | ||
| const torch::Tensor& layout_range, torch::Tensor& active_ranks, | ||
| std::tuple<std::optional<EventHandle>, std::optional<std::function<void()>>> |
There was a problem hiding this comment.
The new raw-pointer dispatch signature for MooncakeEpBuffer in the header makes the native C++ entrypoint very dependent on Python-side validation. In the implementation (mooncake-ep/src/mooncake_ep_buffer.cpp), we now only assert basic shape properties (e.g., hidden alignment and num_tokens <= num_max_dispatch_tokens_per_rank) and then reinterpret the uint64_t arguments as device pointers without additional checks. All of the richer invariants (dtype, contiguity, dimension consistency between x, topk_idx, and the various packed receive buffers) are enforced in mooncake.mooncake_ep_buffer.Buffer.dispatch before calling into C++.
This is fine as long as every call comes via the Python wrapper, but it leaves the C++ API easy to misuse from any future internal caller: passing a null x_ptr when num_tokens > 0, or inconsistent num_* parameters relative to the buffer layout, would satisfy the current assertions and only fail deep in the CUDA kernels. To keep the raw-pointer interface safer without reintroducing Torch C++ dependencies, it might be worth adding a couple of low-cost guards here, such as asserting that num_tokens == 0 || (x_ptr && topk_idx_ptr) and that the packed receive pointers (packed_recv_x_ptr, packed_recv_count_ptr, packed_recv_src_info_ptr, packed_recv_layout_range_ptr) are non-null, plus optional debug-mode checks that Python never passes obviously inconsistent sizes.
🤖 Generated by Qoder • Fix in Qoder
| @@ -443,6 +455,8 @@ def dispatch( | |||
| raise AssertionError("topk_idx and topk_weights must be None when cached handle is provided") | |||
There was a problem hiding this comment.
The new elastic dispatch path moves a lot of deterministic and prefix-sum bookkeeping into Python, which is a reasonable trade-off for decoupling from Torch C++ but does shift some safety responsibilities:
- When a cached
EPHandleis provided, we now explicitly forbid any configuration wheredo_expandis requested orhandle.do_expandis true (raisingAssertionError("Cached EPHandle currently supports only do_expand=False")). That’s a clear constraint, but it means existing callers that cached handles for expanded layouts will start failing at runtime instead of reusing the layout. If that’s intentional, it might be worth documenting that cached handles are strictly non-expanded and that expansion requires fresh handles. - In the
do_cpu_syncblock, the computation ofnum_recv_tokens_per_expert_listand the slicing ofrecv_x,recv_x_scales,recv_topk_idx,recv_topk_weights, andrecv_src_metadatais now implemented in Python based onexpert_psum_cpuandscaleup_psum_cpu. The original C++ implementation usedEP_HOST_ASSERTguards to validate these prefix sums and protect against negative or out-of-range counts. Here we trust the GPU results and truncate tensors accordingly; if a future kernel bug corrupts these sums, we’d get incorrect counts or misaligned slices without a clear assertion.
Given how central these metadata tensors are to elastic routing, would it make sense to add a couple of lightweight sanity checks on the Python side (e.g., non-negative per-expert counts and actual_num_output_tokens not exceeding the allocated capacity) and perhaps a small test that exercises cached-handle reuse and do_expand=True explicitly? That would help keep the new implementation robust while preserving the Torch-free native interface.
🤖 Generated by Qoder • Fix in Qoder
| @@ -548,29 +676,25 @@ def get_next_combine_buffer(self, handle: object): | |||
| hidden, | |||
There was a problem hiding this comment.
Buffer.get_next_combine_buffer now unconditionally returns the Python-side _fallback_next_combine_buffer in the fallback format:
if (
self._fallback_next_combine_buffer is None
or self._fallback_next_combine_buffer.shape
!= (
num_experts // self.group_size,
num_max_dispatch_tokens_per_rank * self.group_size,
hidden,
)
):
self._fallback_next_combine_buffer = torch.empty(...)
return self._fallback_next_combine_bufferSince MooncakeEpBuffer::get_next_combine_buffer has been removed, the native combine path now relies entirely on the pointer passed in as combined_x_ptr, and this Python buffer is just one possible tensor to use. That’s consistent with the new C++ signature, but it does mean that in fast-path mode get_next_combine_buffer() is no longer a view onto an internal RDMA send buffer; it’s simply a staging tensor that tests and callers can populate before passing it into runtime.combine.
If other parts of the codebase or documentation still describe get_next_combine_buffer as exposing the native combine buffer directly (for "zero-copy" behavior), it may be worth updating that language to reflect the new pointer-based design. Otherwise there’s a risk that future code will assume assigning into get_next_combine_buffer() implicitly updates transport-managed memory, when in reality it only affects the tensor you later pass as combined_x.
🤖 Generated by Qoder • Fix in Qoder
There was a problem hiding this comment.
👋 Review Summary
This is a substantial EP refactor that cleanly decouples the native EP layer from torch C++ APIs, pushes shape/dtype validation into Python, and adds a focused native benchmark. Overall the design direction makes sense and the reported test matrix is strong; most surfaces look well thought through.
🛡️ Key Risks & Issues
- Raw-pointer EP APIs:
MooncakeEpBuffer::dispatch/combinenow operate entirely on raw pointers and integer sizes (mooncake-ep/include/mooncake_ep_buffer.h:113–132, mooncake-ep/src/mooncake_ep_buffer.cpp:177–324, 326–355). This is fine when exclusively driven through the validated Python wrappers, but any future native caller or benchmark that bypasses those checks could pass mis-sized buffers or host pointers and the kernels would happily run on arbitrary memory. The Python side currently enforces strong invariants; it would be safer to explicitly document thatMooncakeEpBufferis intended to be used only via these wrappers and to add a couple of extra sanity assertions in the C++ layer to catch obvious misuse early. - Event semantics and overlap: The new EventHandle and
_wait_native_event_on_current_streamhelper change howEventOverlap.current_stream_wait()behaves, especially when the current stream’s native pointer is 0 (mooncake-wheel/mooncake/mooncake_ep_buffer.py:24–35, mooncake-ep/include/mooncake_ep_event.h:9–41). In that case we now do a host-widesynchronize()instead of a stream-local wait, which can reduce overlap and accidentally serialize unrelated work. Given how central EP event-handling is to hiding communication latency, this semantic change deserves explicit confirmation, documentation, and a regression test. - Elastic cached-handle and hybrid routing: The elastic refactor’s raw-pointer
dispatch/combineplus cached EPHandle (native_handle=True) and hybrid topology support introduce complex state (prefix sums, slot indices, forward metadata, linked lists) that must stay perfectly in sync with the current dispatch configuration (mooncake-ep/include/elastic/mooncake_ep_elastic_buffer.h:68–91, mooncake-ep/src/mooncake_ep_elastic_buffer.cpp:230–364, 366–425; mooncake-wheel/mooncake/mooncake_elastic_buffer.py:437–714). If a cached handle is reused after changing num_experts, num_max_tokens_per_rank, expert_alignment, num_sms, or topology, native kernels will operate on stale state and silently produce wrong routing. The Python wrapper already guards some cases; strengthening config matching and adding targeted tests around cached-handle reuse and hybrid mode would help keep this path robust.
🧪 Verification Advice
- Continue running the existing EP test battery (
test_mooncake_ep.py,test_elastic_buffer.py,test_ep_grid.py) on both CUDA and non-CUDA backends where available (MUSA/MACA), withWITH_EP=ON, to validate the new build and pointer-based interfaces end-to-end. - Add targeted tests for elastic deterministic dispatch (with
deterministic=Trueand non-hybrid topology) and elastic FP8 dispatch, as these new code paths are not obviously exercised today but carry non-trivial complexity. - For the new native legacy benchmark, consider a small harness that deliberately passes invalid configurations (e.g. hidden not divisible by 128, zero tokens, or fallback EP) and asserts that
_benchmark_legacy_bufferraises informative errors rather than hanging or misbehaving. - Explicitly exercise EventOverlap on the default stream and under split SEND/RECV and MACA/MUSA modes, confirming that the new
current_stream_waitsemantics deliver the intended overlap and do not regress behavior compared to the previous torch::Event integration.
💡 Thoughts & Suggestions
- The direction of pushing Tensor-level validation and orchestration into Python while keeping the native EP layer on raw pointers and CUDA events is solid and aligns well with the goal of decoupling from torch C++ APIs. Keeping a short comment block or doc note explaining the expected calling pattern (Python only, validated tensors, stream pointer constraints) would make it easier for future contributors to avoid misusing the raw APIs.
- The new legacy benchmark is a nice addition and should be very useful for tracking native EP performance independent of torch; adding a couple of configuration guards/tests, as mentioned above, would make it a more robust diagnostic tool.
- As the elastic subsystem continues to evolve, it may be worth centralizing some of the EPHandle–config consistency checks so that changes in routing strategy or topology automatically surface as assertion failures instead of subtle misrouting.
🤖 Generated by Qoder • View workflow run
| ~MooncakeEpBuffer() noexcept(false); | ||
|
|
||
| std::tuple<torch::Tensor, std::optional<torch::Tensor>, torch::Tensor, | ||
| torch::Tensor, torch::Tensor, std::optional<EventHandle>, | ||
| std::optional<std::function<void()>>> | ||
| dispatch(const torch::Tensor& x, const torch::Tensor& topk_idx, | ||
| torch::Tensor& active_ranks, int num_max_dispatch_tokens_per_rank, | ||
| int num_experts, int timeout_us, bool use_fp8, bool async, | ||
| bool return_recv_hook); | ||
|
|
||
| std::tuple<torch::Tensor, std::optional<EventHandle>, | ||
| std::optional<std::function<void()>>> | ||
| combine(const torch::Tensor& x, const torch::Tensor& topk_idx, | ||
| const torch::Tensor& topk_weights, const torch::Tensor& src_info, | ||
| const torch::Tensor& layout_range, torch::Tensor& active_ranks, | ||
| std::tuple<std::optional<EventHandle>, std::optional<std::function<void()>>> |
There was a problem hiding this comment.
The raw-pointer refactor for MooncakeEpBuffer::dispatch/combine looks clean overall, but it now relies almost entirely on Python-side invariants for safety.
In the new signatures (mooncake-ep/include/mooncake_ep_buffer.h:113–132, implemented in mooncake-ep/src/mooncake_ep_buffer.cpp:177–324 and 326–355), we only assert expert divisibility, hidden alignment, token upper bounds, and non-null pointers. We no longer have access to Tensor metadata like strides, device, or contiguity, so any caller that bypasses mooncake_ep_buffer.py and passes mis-sized buffers, host pointers, or mismatched hidden vs allocated width will result in the kernels operating on arbitrary memory.
Python wrappers currently enforce strong shape/dtype checks, but the new benchmark and any future native users of MooncakeEpBuffer will be able to plug in arbitrary pointers. To keep this refactor robust, it would be helpful to:
- Explicitly state in comments/docs that
MooncakeEpBufferis intended to be called only through the validated Python wrappers, and - Consider adding a few extra sanity assertions here (e.g.,
num_tokens == 0implies all payload pointers are either null or point to at least one element,packed_recv_x_scales_ptrnon-null only whenuse_fp8is true, etc.) to catch obvious misuse early.
This is especially important for maintainers who might later add a C++-side caller and assume the API still behaves like the old torch::Tensor version.
🤖 Generated by Qoder • Fix in Qoder
| std::tuple<int, int> get_logical_domain_size() const; | ||
| int get_theoretical_num_sms(int num_experts, int num_topk) const; | ||
|
|
||
| ElasticDispatchOutput dispatch( | ||
| const torch::Tensor& x, const std::optional<torch::Tensor>& sf, | ||
| const torch::Tensor& topk_idx, | ||
| const std::optional<torch::Tensor>& topk_weights, | ||
| torch::Tensor& active_ranks, int num_experts, | ||
| int num_max_tokens_per_rank, int expert_alignment, int num_sms, | ||
| bool do_expand, bool do_cpu_sync, bool async_with_compute_stream, | ||
| const std::optional<ElasticNativeHandle>& cached_handle = std::nullopt); | ||
|
|
||
| ElasticCombineOutput combine( | ||
| const torch::Tensor& x, const ElasticNativeHandle& handle, | ||
| const std::optional<torch::Tensor>& topk_weights, | ||
| torch::Tensor& active_ranks, int num_sms, | ||
| bool async_with_compute_stream, | ||
| const std::optional<torch::Tensor>& out); | ||
| std::optional<EventHandle> dispatch( |
There was a problem hiding this comment.
The elastic refactor introduces raw-pointer dispatch/combine APIs and a deterministic cached-handle path that look quite powerful but also fairly fragile under configuration mismatches.
In MooncakeElasticBuffer::dispatch (mooncake-ep/include/elastic/mooncake_ep_elastic_buffer.h:68–81, mooncake-ep/src/mooncake_ep_elastic_buffer.cpp:230–364), the kernel now receives pointers to preallocated prefix-sum buffers, slot indices, forward metadata, and channel linked lists, with a cached_mode flag. The Python wrapper in mooncake_elastic_buffer.py rebuilds or reuses these tensors based on an EPHandle and passes them as raw pointers.
The native code assumes that the cached handle’s configuration (num_experts, num_max_tokens_per_rank, expert_alignment, num_sms, hybrid topology, do_expand) stays perfectly in sync with the current dispatch settings. If a caller accidentally reuses a cached handle after changing any of these knobs, the kernels will operate on stale slot maps and prefix sums, leading to misrouted tokens and corrupted recv_src_metadata that are hard to debug.
The Python side already guards some cases (e.g. rejecting cached handles when do_expand is requested), but this path is complex enough that it would be good to:
- Strengthen validation either in Python or C++ to assert that the EPHandle configuration exactly matches the current dispatch parameters when
cached_mode=True, and - Add targeted tests that reuse cached handles across sequential dispatches while varying num_experts/num_max_tokens_per_rank/topology to ensure we either rebuild the handle or fail fast rather than silently producing incorrect routing.
This will help keep the cached-handle feature safe as the elastic EP evolves.
🤖 Generated by Qoder • Fix in Qoder
| def _wait_native_event_on_current_stream(event: "ep.EventHandle") -> None: | ||
| stream_ptr = _native_current_stream_ptr() | ||
| if stream_ptr == 0: | ||
| event.synchronize() | ||
| return | ||
| event.current_stream_wait(stream_ptr) | ||
|
|
||
|
|
||
| class EventOverlap: |
There was a problem hiding this comment.
The new _wait_native_event_on_current_stream helper changes EventOverlap.current_stream_wait() semantics when the current stream maps to a native pointer of 0.
For non-zero streams, we call event.current_stream_wait(stream_ptr) (mooncake-wheel/mooncake/mooncake_ep_buffer.py:28–33), which uses cudaStreamWaitEvent(stream, *event, 0) under the hood (mooncake-ep/include/mooncake_ep_event.h:36–39). When torch.cuda.current_stream().cuda_stream is 0, we instead call event.synchronize(), which blocks the host until the event completes.
That special case means that code running on a default/native stream with pointer 0 will serialize all work rather than establishing a stream-local dependency, potentially reducing overlap between EP communication and user compute compared to the previous torch::Event behavior.
Given how central EventOverlap is to hiding EP latency, can we either confirm this global-sync semantics is intentional (and document it for CUDA/MUSA/MACA), or adjust the implementation so that default-stream usage still behaves like a stream wait where possible? A small regression test that exercises an EP call on the default stream and asserts the expected waiting behavior would also help catch regressions here.
🤖 Generated by Qoder • Fix in Qoder
There was a problem hiding this comment.
👋 Review Summary
This refactor cleanly decouples the EP native layer from Torch C++ APIs, moves tensor/stream management into the Python wrappers, and adds a focused legacy benchmark. The overall direction looks solid and the described test runs cover the main EP and ElasticBuffer flows across realistic configurations.
🛡️ Key Risks & Issues
- EP Python fallback path: EventOverlap now routes through
_wait_native_event_on_current_stream, assuming anEventHandlewithcurrent_stream_wait(stream_ptr). The_DummyEventused by the fallback path still exposes a no-argcurrent_stream_wait, so when native EP is unavailable or_use_fallbackis true, async EP operations can raise a TypeError instead of synchronizing. This affects robustness precisely in the environments that rely on the fallback. - Native EP combine buffer optimization:
Buffer.get_next_combine_bufferwas simplified to always return a CUDA bfloat16 buffer and no longer wraps the native RDMA send buffer. On the fast path this likely reintroduces an extra copy and narrows the performance gap between fast and fallback modes, which is worth revisiting if the API is still intended as a transport-backed optimization. - Test-grid port selection: The new
_pick_free_porthelper fixes hard-coded 29500 collisions within the grid suite, but it still leaves a race window where another process can bind the chosen port between closing the temporary socket and initializing the process group. On busy shared hosts that can produce intermittent bootstrap failures that don’t obviously point back to EP.
🧪 Verification Advice
- Explicitly exercise the EP Python fallback path (forcing
_use_fallback/ disabling native_ep) with async dispatch/combine and confirm that EventOverlap works without TypeErrors and still synchronizes GPU work. Include bothcurrent_stream_waitandsynchronize()on the dummy event. - Run the existing EP and ElasticBuffer tests with additional assertions around invalid shapes/dtypes, mismatched
num_experts % group_size, and inconsistentactive_ranks, to lock in the new Python-side validation invariants. - For the legacy benchmark, consider adding a lightweight CI test that uses a very small configuration and asserts that invalid configs raise the expected exceptions while a valid config completes and produces non-degenerate timing data.
- In the grid tests, try simulating port contention (or mocking port selection) to confirm behavior under conflicts and, if you add retry or clearer failure messaging, validate those paths explicitly.
💡 Thoughts & Suggestions
- The pointer-based APIs in
MooncakeEpBufferandMooncakeElasticBufferlook consistent with the Python wrappers, but they are more fragile than the previous tensor-based interfaces. It will be helpful to keep strong invariants documented and backed by tests, especially around hybrid-mode metadata shapes and cached-handle usage, so future changes don’t accidentally violate those contracts. - If the public EventHandle API is meant to be used directly outside these wrappers, consider a small compatibility shim or helper so existing call sites that relied on
EventHandle()+current_stream_wait()without arguments fail in a controlled, well-documented way (or remain supported) rather than surprising users with argument errors. - Overall, the decoupling and build changes move EP toward a more self-contained native layer and should pay off in long-term maintainability; tightening the few edge cases called out above would help keep the behavior robust across both native and fallback deployments.
🤖 Generated by Qoder • View workflow run
| @@ -579,6 +703,9 @@ class _DummyEvent: | |||
| def current_stream_wait(self): | |||
There was a problem hiding this comment.
In the Buffer Python wrapper (mooncake-wheel/mooncake/mooncake_ep_buffer.py), _wait_native_event_on_current_stream unconditionally calls event.current_stream_wait(stream_ptr), which matches the new native ep.EventHandle signature but not the _DummyEvent used on the EP fallback path.
_DummyEvent.current_stream_wait still takes no arguments, so whenever the runtime falls back to _DummyEvent (for example when _use_fallback is true or native _ep is unavailable) and EventOverlap tries to wait on the event, Python will raise a TypeError instead of synchronizing work. That would break the fallback path exactly in the scenarios where robustness matters most.
To keep the fallback behavior safe, we should either update _DummyEvent.current_stream_wait to accept an unused stream_ptr parameter or teach _wait_native_event_on_current_stream to detect the dummy event and call the no-arg method. That way both native and fallback events remain compatible with the EventOverlap helper, and async dispatch/combine behavior stays consistent across EP implementations.
🤖 Generated by Qoder • Fix in Qoder
| @@ -548,29 +676,25 @@ def get_next_combine_buffer(self, handle: object): | |||
| hidden, | |||
There was a problem hiding this comment.
The Buffer.get_next_combine_buffer implementation now always returns a CUDA bfloat16 tensor stored in _fallback_next_combine_buffer, and the branch that previously delegated to self.runtime.get_next_combine_buffer on the native fast path has been removed.
On the native EP path, the older implementation wrapped the RDMA send buffer directly as a Tensor, allowing the combine kernel to write into transport-managed memory without an extra staging allocation. With the new version, callers always receive a separate heap-allocated buffer, which likely forces an additional copy and reduces the performance benefit of the fast path for large expert/dispatch configurations.
If we still intend to expose a "next combine buffer" optimized path for native EP, it might be worth preserving a runtime-backed variant (e.g., only using _fallback_next_combine_buffer when _use_fallback is true) so that fast-path users can keep writing directly into the RDMA send buffer rather than an intermediate tensor.
🤖 Generated by Qoder • Fix in Qoder
| def setUp(self): | ||
| import_torchada_if_needed() | ||
| self.world_size = torch.cuda.device_count() | ||
| os.environ["MASTER_ADDR"] = "127.0.0.1" |
There was a problem hiding this comment.
In TestMooncakeEPBuffer.setUp (mooncake-ep/tests/test_ep_grid.py), the grid tests now use _pick_free_port() to bind an ephemeral TCP port and then set MASTER_PORT after closing the temporary socket. This avoids hard-coded 29500 collisions between fail-rank cases, but it doesn’t fully eliminate races with other processes on the same host.
Between closing the temporary socket and calling torch.distributed.init_process_group, another process can bind the same port, leading to intermittent bootstrap failures or noisy connection errors that are hard to attribute to the grid tests themselves.
If we want to make these tests more robust under contention, we might consider retrying on bind failure with a small number of fallback ports or keeping the socket open until after the process group is initialized so that the chosen port cannot be stolen between selection and bootstrap.
🤖 Generated by Qoder • Fix in Qoder
|
The majority of CI have passed (https://github.com/kvcache-ai/Mooncake/actions/runs/29676976479?pr=2883). Also tested e2e on SGLang and observed no regression. This should significantly improve Mooncake's internal dev experience and reduce CI build time. cc: @yuechen-sys |
2 participants
Description
This PR removes the Mooncake EP native dependency on torch C++ APIs and builds the EP extension directly from the main Mooncake CMake graph.
The native EP layer now uses raw pointers, native stream handles, and native CUDA events instead of
torch::Tensor,torch::Event, andat::cuda::*. Tensor allocation, validation, stream lookup, and output slicing move to the Python wrappers for both legacyBufferand elasticElasticBuffer.The EP build no longer uses
mooncake-ep/setup.pyortorch.utils.cpp_extension. The top-levelWITH_EPbuild now producesmooncake._epdirectly viapybind11_add_module.This PR also adds a private C++ legacy EP benchmark. Python performs distributed bootstrap and owns the fixed tensors, while the timed dispatch/combine loop uses the native raw-pointer API. This measures native EP performance without reintroducing Torch C++ dependencies or expanding the public EP API.
A follow-up validation run on
test_ep_grid.pyalso fixed two EP-adjacent issues on the same branch: the Python dummy-event compatibility path now implementssynchronize(), and the grid test now allocates a freshMASTER_PORTper case so fail-rank cases do not collide on29500.Module
mooncake-ep)mooncake-integration)mooncake-wheel)mooncake-transfer-engine)mooncake-store)mooncake-pg)mooncake-p2p-store)mooncake-common)mooncake-rl)Type of Change
How Has This Been Tested?
Core EP behavior was validated on
h20/qjh003in containerep-torch-decouplingwithMC_TE_FILTERS=mlx5_1. The native benchmark andSGLang E2E comparison ran on
h20/qjh002with the same explicit HCA filter.Test commands:
Test results:
Manual testing:
test_mooncake_ep.pypassed with refactoredmooncake._eptest_elastic_buffer.pypassed and printedMOONCAKE_ELASTIC_TEST_OK world=2 route=local recv=32 expanded=256 scaleout=1 scaleup=2test_ep_grid.pypassed withRan 49 tests in 210.747sandOKmlx5_1:119.10-119.27 us, about246 GB/sfor dispatch + combine.163.12s/163.34-163.38s(0.14% delta); the 200-query GSM8K run passed for both configurations (11.449s/11.654s,1586.528/1517.910 token/s). No functional or CUDA-graph capture regression was observed.Checklist
./scripts/code_format.shpre-commit run --all-filesand all hooks passAI Assistance Disclosure
Used Codex/GPT-5 as an implementation assistant for code restructuring, build/test iteration, and review-log drafting. All code paths and test results were manually inspected and validated by the submitter.