[GH-ISSUE #3350] [RFC]: Decouple native EP from the PyTorch C++ API #850

New issue

Closed

opened 2026-10-29 21:20:11 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

2026-10-29 21:20:11 +00:00

Owner

Originally created by @UNIDY2002 on GitHub (Aug 10, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/3350

Changes proposed

Motivation

Mooncake EP currently builds its native extension through PyTorch C++ APIs and a setup.py-driven staging flow. The native implementation takes torch::Tensor and ATen CUDA stream objects directly, which ties the EP binary to a specific PyTorch C++ ABI. Release builds consequently compile and stage an EP artifact for every supported PyTorch version, even though the EP communication and kernel implementation does not require PyTorch at runtime.

This RFC establishes a native EP boundary based on device pointers, tensor metadata, CUDA-compatible streams, and CUDA-compatible events. The Python wheel owns tensor validation, allocation, stream-handle extraction, and lifetime management. The native library owns transport setup, workspace management, kernel launches, and event synchronization.

The new boundary removes the per-PyTorch-version native EP build matrix. Releases still build the required Python ABI artifacts, while one native EP implementation serves supported PyTorch versions. This shortens release builds, removes the EP-specific setup.py staging path, and makes the build graph easier to reason about and maintain.

Proposed changes

Build mooncake-ep as a first-class CMake target and remove its setup.py and BuildEpExt.cmake build/staging flow.
Replace torch::Tensor and ATen CUDA stream types in the native legacy and elastic EP APIs with device pointers, shape and dtype metadata, and native stream handles.
Move tensor validation, output allocation, stream extraction, and async tensor lifetime retention into the Python wheel wrappers.
Replace Torch event ownership in the native layer with CUDA-compatible events and explicit keepalive ownership for asynchronous elastic dispatch scratch.
Preserve the current Python Buffer and ElasticBuffer behavior while routing them through the new native ABI.
Add a native legacy-buffer benchmark in C++ with a Python launcher. It provides a stable correctness and latency regression harness independent of the Python-level EP wrapper.

Scope and compatibility

The RFC covers both legacy and elastic EP. They share the same native build target and ABI boundary, and migrating them together removes the existing split build path.

The implementation retains CUDA behavior and adds the build/runtime compatibility required for MUSA through the same native boundary. MUSA protocol details and platform-specific performance tuning are implementation follow-ups; they are not the motivation for this RFC.

The public Python EP API remains unchanged. Existing callers continue to construct Buffer and ElasticBuffer and invoke dispatch and combine with tensors.

Validation plan

Build the native EP target with the supported CUDA and MUSA configurations.
Run legacy EP correctness grids for P2P and IBGDA transports.
Run elastic EP correctness coverage through the existing Python wrapper.
Run the new native legacy-buffer benchmark as a release regression gate.
Compare CUDA release artifacts across supported PyTorch versions to verify that EP no longer requires a separate native build for each PyTorch version.

Before submitting a new issue...

Searched existing issues for related EP build and PyTorch ABI discussions.
Read the Mooncake documentation and confirmed the RFC scope.

Originally created by @UNIDY2002 on GitHub (Aug 10, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/3350 ## Changes proposed ### Motivation Mooncake EP currently builds its native extension through PyTorch C++ APIs and a `setup.py`-driven staging flow. The native implementation takes `torch::Tensor` and ATen CUDA stream objects directly, which ties the EP binary to a specific PyTorch C++ ABI. Release builds consequently compile and stage an EP artifact for every supported PyTorch version, even though the EP communication and kernel implementation does not require PyTorch at runtime. This RFC establishes a native EP boundary based on device pointers, tensor metadata, CUDA-compatible streams, and CUDA-compatible events. The Python wheel owns tensor validation, allocation, stream-handle extraction, and lifetime management. The native library owns transport setup, workspace management, kernel launches, and event synchronization. The new boundary removes the per-PyTorch-version native EP build matrix. Releases still build the required Python ABI artifacts, while one native EP implementation serves supported PyTorch versions. This shortens release builds, removes the EP-specific `setup.py` staging path, and makes the build graph easier to reason about and maintain. ### Proposed changes - Build `mooncake-ep` as a first-class CMake target and remove its `setup.py` and `BuildEpExt.cmake` build/staging flow. - Replace `torch::Tensor` and ATen CUDA stream types in the native legacy and elastic EP APIs with device pointers, shape and dtype metadata, and native stream handles. - Move tensor validation, output allocation, stream extraction, and async tensor lifetime retention into the Python wheel wrappers. - Replace Torch event ownership in the native layer with CUDA-compatible events and explicit keepalive ownership for asynchronous elastic dispatch scratch. - Preserve the current Python `Buffer` and `ElasticBuffer` behavior while routing them through the new native ABI. - Add a native legacy-buffer benchmark in C++ with a Python launcher. It provides a stable correctness and latency regression harness independent of the Python-level EP wrapper. ### Scope and compatibility The RFC covers both legacy and elastic EP. They share the same native build target and ABI boundary, and migrating them together removes the existing split build path. The implementation retains CUDA behavior and adds the build/runtime compatibility required for MUSA through the same native boundary. MUSA protocol details and platform-specific performance tuning are implementation follow-ups; they are not the motivation for this RFC. The public Python EP API remains unchanged. Existing callers continue to construct `Buffer` and `ElasticBuffer` and invoke `dispatch` and `combine` with tensors. ### Validation plan - Build the native EP target with the supported CUDA and MUSA configurations. - Run legacy EP correctness grids for P2P and IBGDA transports. - Run elastic EP correctness coverage through the existing Python wrapper. - Run the new native legacy-buffer benchmark as a release regression gate. - Compare CUDA release artifacts across supported PyTorch versions to verify that EP no longer requires a separate native build for each PyTorch version. ## Before submitting a new issue... - [x] Searched existing issues for related EP build and PyTorch ABI discussions. - [x] Read the Mooncake documentation and confirmed the RFC scope.

bsakaguchi

2026-10-29 21:20:11 +00:00

closed this issue
added the
RFC
label

bsakaguchi commented

2026-10-29 21:20:12 +00:00

Author

Owner

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

Thanks for opening this issue, @UNIDY2002!

Field	Value
Issue	#3350
GitHub user ID	`41641284`
Reporter	@UNIDY2002

A maintainer will triage this when possible. To help us respond faster, please include:

Mooncake version or commit SHA
Environment (OS, CUDA/driver, RDMA stack if relevant)
Steps to reproduce and expected vs. actual behavior

Useful links: Documentation · Contributing guide

This message was posted automatically by the issue bot.

@github-actions[bot] commented on GitHub (Aug 10, 2026): Thanks for opening this issue, @UNIDY2002! | Field | Value | |-------|-------| | **Issue** | #3350 | | **GitHub user ID** | `41641284` | | **Reporter** | @UNIDY2002 | A maintainer will triage this when possible. To help us respond faster, please include: - Mooncake version or commit SHA - Environment (OS, CUDA/driver, RDMA stack if relevant) - Steps to reproduce and expected vs. actual behavior Useful links: [Documentation](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/kvcache-ai/Mooncake/blob/main/CONTRIBUTING.md) > This message was posted automatically by the issue bot.