kvcache-ai/Mooncake

Fork 0

mirror of https://github.com/kvcache-ai/Mooncake.git synced 2026-11-05 01:17:55 +00:00

[GH-ISSUE #3200] [RFC]: Decouple Mooncake PG from PyTorch #819

New issue

Closed

opened 2026-10-29 21:19:52 +00:00 by bsakaguchi · 2 comments

bsakaguchi commented

2026-10-29 21:19:52 +00:00

Owner

Originally created by @caozhanhao on GitHub (Jul 30, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/3200

Originally assigned to: @caozhanhao on GitHub.

Summary

This RFC proposes decoupling the Mooncake PG into two distinct layers:

PG Core (libmooncake_pg.so): A separate C++ shared library built alongside the Transfer Engine via the top-level CMake project. It contains zero dependencies on PyTorch.
PG Torch: A lightweight PyTorch extension that maintains the existing c10d and Python interfaces by acting as an adapter on top of PG Core.

PG Core will expose a preliminary C API (mooncake_pg.h) for native C++ tests, sanitizers, micro-benchmarks, and non-PyTorch runtime integrations. The initial API is experimental and does not guarantee source or ABI compatibility.

Motivation

Currently, Mooncake PG is compiled as a PyTorch extension, whereas the Transfer Engine is compiled directly by the top-level CMake build system. Because these two components are built under different build contexts, shared header-only dependencies can be selected or configured inconsistently across the build boundary. This leads to incompatible template instantiations and subtle ABI mismatches at runtime (such as the YLT header mismatch addressed in #3164).

Furthermore, the existing build architecture tightly couples the entire PG implementation to the PyTorch C++ ABI. Consequently:

ABI Fragmentation: Every supported PyTorch C++ ABI version requires a complete, dedicated rebuild of the entire PG library.
Heavyweight Testing Overhead: Native sanitizer runs and standalone tests must link against PyTorch, while the existing test workflow also runs through Python.

By migrating the core communication logic into the top-level CMake build, PG Core and the Transfer Engine will share a single, unified dependency graph. PyTorch ABI-sensitive code is strictly isolated to the thin PG Torch adapter layer.

Architecture & Design

+-----------------------+
|  Native Application   |
+-----------+-----------+
            |
            | (mooncake_pg.h)
            v
+-----------------------+              +-----------------------+
|        PG Core        | -----------> |    Transfer Engine    |
| (`libmooncake_pg.so`) |              +-----------------------+
+-----------------------+
            ^
            |
            | (C API Boundary)
+-----------+-----------+
|       PG Torch        |
+-----------+-----------+
            |
            v
+-----------------------+
|   torch.distributed   |
+-----------------------+

Layer Boundary (`mooncake_pg.h`)

PG Core: Accepts opaque handles, raw memory buffers, element counts, data types, rank IDs, and GPU streams. It manages the Coordinator and Agent state machines, communicator membership, topology, collective and point-to-point (P2P) algorithms, fault tolerance, and asynchronous completion tracking. It contains no at::Tensor, c10d classes, PyTorch Store, or Python C API bindings.
PG Torch: Functions as a thin adapter layer. It unpacks PyTorch at::Tensor objects, extracts device pointers (data_ptr()) and active stream handles, maps c10d options to core C enums, invokes the PG Core C API, and wraps core completion handles or GPU stream events into c10d::Work instances. It does not re-implement any control plane or collective algorithms.

Key Abstractions

MooncakeCommunicator: The primary C++ class inside PG Core that encapsulates the state for a single CPU or GPU communication group (rank mappings, device context, buffers, worker handles, and membership view). The C API exposes this class as an opaque handle: mooncakePgComm_t.
MooncakeBackend: The c10d::ProcessGroup implementation remaining inside PG Torch, serving as the adapter for an underlying mooncakePgComm_t handle.

Build System Integration

The top-level CMake build compiles libmooncake_pg.so alongside the Transfer Engine. PG Torch is built as a lightweight PyTorch extension once per target PyTorch ABI, linking dynamically against libmooncake_pg.so.

Why a C API?

PG Core and PG Torch are compiled by different toolchains. In particular, PG Torch must follow PyTorch's _GLIBCXX_USE_CXX11_ABI setting, which may differ from the ABI used by the top-level CMake build. Exposing C++ types such as std::string, std::vector, exceptions, or class layouts across this boundary would couple PG Core to each PyTorch C++ ABI and require rebuilding it for different PyTorch distributions.

The C API acts as an ABI boundary by exposing only opaque handles and C-compatible values. MooncakeCommunicator remains the internal C++ implementation, while PG Torch and other external users call it through mooncake_pg.h. The initial C API may still evolve, but it prevents PyTorch's C++ ABI from leaking into PG Core.

C API

The preliminary C API header lives at mooncake-pg/include/mooncake_pg.h.

Types

Opaque Handles: mooncakePgContext_t, mooncakePgComm_t, mooncakePgCompletion_t, mooncakePgStream_t
Error Handling: mooncakePgResult_t
Data Types & Operations: mooncakePgDataType_t, mooncakePgReduceOp_t
Configuration: mooncakePgDeviceType_t, mooncakePgIdResolvePolicy_t, mooncakePgCommConfig_t
Elasticity: mooncakePgProposalStatus_t, mooncakePgProposalResponse_t, mooncakePgSyncAfterFailureStatus_t, mooncakePgSyncAfterFailureResponse_t

Functions

Error Handling

mooncakePgGetErrorString
mooncakePgGetLastError

Context Lifecycle & Global Options

mooncakePgContextCreate / mooncakePgContextDestroy
mooncakePgContextInitialize
mooncakePgContextLaunchCoordinator / mooncakePgContextConnectCoordinator
mooncakePgContextSetHostIp
mooncakePgContextSetTransferEngine
mooncakePgContextSetDeviceFilter
mooncakePgContextSetCollectiveTimeout / mooncakePgContextSetP2PTimeout
mooncakePgContextSetFaultReconciliationWindow

Communicator Management

mooncakePgCommCreate / mooncakePgCommDestroy
mooncakePgCommGetRank / mooncakePgCommGetSize / mooncakePgCommGetMaxGroupSize

Collective Operations

Each collective provides separate CPU and GPU entry points, named mooncakePg<Operation>Cpu and mooncakePg<Operation>Gpu, respectively:

Broadcast, AllReduce, AllGather, ReduceScatter, AllToAll, Reduce, Gather, Scatter, Barrier

Point-to-Point Communication

mooncakePgSendCpu / mooncakePgRecvCpu
mooncakePgSendGpu / mooncakePgRecvGpu

Asynchronous Completion Tracking

mooncakePgCompletionIsCompleted
mooncakePgCompletionWait
mooncakePgCompletionDestroy

Semantics: CPU collectives and all P2P operations return an asynchronous mooncakePgCompletion_t handle. GPU collectives execute asynchronously on the caller-provided hardware stream without returning a completion handle. Callers are responsible for maintaining memory buffer validity and stream lifecycle until work execution completes on the hardware.

Elastic Membership & Topology

mooncakePgCommGetActiveRanks
mooncakePgCommGetPeerState
mooncakePgCommActivateRanks / mooncakePgCommDeactivateRanks
mooncakePgCommJoin
mooncakePgCommSyncAfterFailure
mooncakePgCommGetEpoch
mooncakePgCommGetNumSyncedRanks
mooncakePgCommGetPreferredHca

Alternatives Considered

Move directly to a Python-level ProcessGroup (PyProcessGroup).
This approach would eliminate the C++ dependency on PyTorch entirely. However, PyProcessGroup is a newer and less mature integration path. Changing both the core boundary and the PyTorch integration model at the same time would introduce unnecessary compatibility risk.

The proposed split preserves the established C++ c10d::ProcessGroup integration while removing PyTorch from PG Core. It also leaves open the option of migrating to PyProcessGroup once that integration path is sufficiently mature across the supported PyTorch versions.

Originally created by @caozhanhao on GitHub (Jul 30, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/3200 Originally assigned to: @caozhanhao on GitHub. ## Summary This RFC proposes decoupling the Mooncake PG into two distinct layers: * **PG Core (`libmooncake_pg.so`)**: A separate C++ shared library built alongside the Transfer Engine via the top-level CMake project. It contains zero dependencies on PyTorch. * **PG Torch**: A lightweight PyTorch extension that maintains the existing `c10d` and Python interfaces by acting as an adapter on top of PG Core. PG Core will expose a preliminary C API (`mooncake_pg.h`) for native C++ tests, sanitizers, micro-benchmarks, and non-PyTorch runtime integrations. The initial API is experimental and does not guarantee source or ABI compatibility. ## Motivation Currently, Mooncake PG is compiled as a PyTorch extension, whereas the Transfer Engine is compiled directly by the top-level CMake build system. Because these two components are built under different build contexts, shared header-only dependencies can be selected or configured inconsistently across the build boundary. This leads to incompatible template instantiations and subtle ABI mismatches at runtime (such as the YLT header mismatch addressed in #3164). Furthermore, the existing build architecture tightly couples the entire PG implementation to the PyTorch C++ ABI. Consequently: 1. **ABI Fragmentation**: Every supported PyTorch C++ ABI version requires a complete, dedicated rebuild of the entire PG library. 2. **Heavyweight Testing Overhead**: Native sanitizer runs and standalone tests must link against PyTorch, while the existing test workflow also runs through Python. By migrating the core communication logic into the top-level CMake build, PG Core and the Transfer Engine will share a single, unified dependency graph. PyTorch ABI-sensitive code is strictly isolated to the thin `PG Torch` adapter layer. ## Architecture & Design ```text +-----------------------+ | Native Application | +-----------+-----------+ | | (mooncake_pg.h) v +-----------------------+ +-----------------------+ | PG Core | -----------> | Transfer Engine | | (`libmooncake_pg.so`) | +-----------------------+ +-----------------------+ ^ | | (C API Boundary) +-----------+-----------+ | PG Torch | +-----------+-----------+ | v +-----------------------+ | torch.distributed | +-----------------------+ ``` ### Layer Boundary (`mooncake_pg.h`) * **PG Core**: Accepts opaque handles, raw memory buffers, element counts, data types, rank IDs, and GPU streams. It manages the Coordinator and Agent state machines, communicator membership, topology, collective and point-to-point (P2P) algorithms, fault tolerance, and asynchronous completion tracking. It contains no `at::Tensor`, `c10d` classes, PyTorch `Store`, or Python C API bindings. * **PG Torch**: Functions as a thin adapter layer. It unpacks PyTorch `at::Tensor` objects, extracts device pointers (`data_ptr()`) and active stream handles, maps `c10d` options to core C enums, invokes the PG Core C API, and wraps core completion handles or GPU stream events into `c10d::Work` instances. It does not re-implement any control plane or collective algorithms. ### Key Abstractions * **`MooncakeCommunicator`**: The primary C++ class inside PG Core that encapsulates the state for a single CPU or GPU communication group (rank mappings, device context, buffers, worker handles, and membership view). The C API exposes this class as an opaque handle: `mooncakePgComm_t`. * **`MooncakeBackend`**: The `c10d::ProcessGroup` implementation remaining inside PG Torch, serving as the adapter for an underlying `mooncakePgComm_t` handle. ### Build System Integration The top-level CMake build compiles `libmooncake_pg.so` alongside the Transfer Engine. PG Torch is built as a lightweight PyTorch extension once per target PyTorch ABI, linking dynamically against `libmooncake_pg.so`. ### Why a C API? PG Core and PG Torch are compiled by different toolchains. In particular, PG Torch must follow PyTorch's `_GLIBCXX_USE_CXX11_ABI` setting, which may differ from the ABI used by the top-level CMake build. Exposing C++ types such as `std::string`, `std::vector`, exceptions, or class layouts across this boundary would couple PG Core to each PyTorch C++ ABI and require rebuilding it for different PyTorch distributions. The C API acts as an ABI boundary by exposing only opaque handles and C-compatible values. `MooncakeCommunicator` remains the internal C++ implementation, while PG Torch and other external users call it through `mooncake_pg.h`. The initial C API may still evolve, but it prevents PyTorch's C++ ABI from leaking into PG Core. ## C API The preliminary C API header lives at `mooncake-pg/include/mooncake_pg.h`. ### Types * **Opaque Handles**: `mooncakePgContext_t`, `mooncakePgComm_t`, `mooncakePgCompletion_t`, `mooncakePgStream_t` * **Error Handling**: `mooncakePgResult_t` * **Data Types & Operations**: `mooncakePgDataType_t`, `mooncakePgReduceOp_t` * **Configuration**: `mooncakePgDeviceType_t`, `mooncakePgIdResolvePolicy_t`, `mooncakePgCommConfig_t` * **Elasticity**: `mooncakePgProposalStatus_t`, `mooncakePgProposalResponse_t`, `mooncakePgSyncAfterFailureStatus_t`, `mooncakePgSyncAfterFailureResponse_t` ### Functions #### Error Handling * `mooncakePgGetErrorString` * `mooncakePgGetLastError` #### Context Lifecycle & Global Options * `mooncakePgContextCreate` / `mooncakePgContextDestroy` * `mooncakePgContextInitialize` * `mooncakePgContextLaunchCoordinator` / `mooncakePgContextConnectCoordinator` * `mooncakePgContextSetHostIp` * `mooncakePgContextSetTransferEngine` * `mooncakePgContextSetDeviceFilter` * `mooncakePgContextSetCollectiveTimeout` / `mooncakePgContextSetP2PTimeout` * `mooncakePgContextSetFaultReconciliationWindow` #### Communicator Management * `mooncakePgCommCreate` / `mooncakePgCommDestroy` * `mooncakePgCommGetRank` / `mooncakePgCommGetSize` / `mooncakePgCommGetMaxGroupSize` #### Collective Operations Each collective provides separate CPU and GPU entry points, named `mooncakePg<Operation>Cpu` and `mooncakePg<Operation>Gpu`, respectively: * `Broadcast`, `AllReduce`, `AllGather`, `ReduceScatter`, `AllToAll`, `Reduce`, `Gather`, `Scatter`, `Barrier` #### Point-to-Point Communication * `mooncakePgSendCpu` / `mooncakePgRecvCpu` * `mooncakePgSendGpu` / `mooncakePgRecvGpu` #### Asynchronous Completion Tracking * `mooncakePgCompletionIsCompleted` * `mooncakePgCompletionWait` * `mooncakePgCompletionDestroy` *Semantics*: CPU collectives and all P2P operations return an asynchronous `mooncakePgCompletion_t` handle. GPU collectives execute asynchronously on the caller-provided hardware stream without returning a completion handle. Callers are responsible for maintaining memory buffer validity and stream lifecycle until work execution completes on the hardware. #### Elastic Membership & Topology * `mooncakePgCommGetActiveRanks` * `mooncakePgCommGetPeerState` * `mooncakePgCommActivateRanks` / `mooncakePgCommDeactivateRanks` * `mooncakePgCommJoin` * `mooncakePgCommSyncAfterFailure` * `mooncakePgCommGetEpoch` * `mooncakePgCommGetNumSyncedRanks` * `mooncakePgCommGetPreferredHca` ## Alternatives Considered - Move directly to a Python-level ProcessGroup (`PyProcessGroup`). This approach would eliminate the C++ dependency on PyTorch entirely. However, `PyProcessGroup` is a newer and less mature integration path. Changing both the core boundary and the PyTorch integration model at the same time would introduce unnecessary compatibility risk. The proposed split preserves the established C++ `c10d::ProcessGroup` integration while removing PyTorch from PG Core. It also leaves open the option of migrating to `PyProcessGroup` once that integration path is sufficiently mature across the supported PyTorch versions.

bsakaguchi

2026-10-29 21:19:52 +00:00

closed this issue
added the
RFC

PyTorch Backend
labels

bsakaguchi commented

2026-10-29 21:19:53 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Jul 30, 2026):

Thanks for opening this issue, @caozhanhao!

Field	Value
Issue	#3200
GitHub user ID	`64004730`
Reporter	@caozhanhao

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 (Jul 30, 2026): Thanks for opening this issue, @caozhanhao! | Field | Value | |-------|-------| | **Issue** | #3200 | | **GitHub user ID** | `64004730` | | **Reporter** | @caozhanhao | 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.

bsakaguchi commented

2026-10-29 21:19:53 +00:00

Author

Owner

@UNIDY2002 commented on GitHub (Jul 31, 2026):

LGTM

@UNIDY2002 commented on GitHub (Jul 31, 2026): LGTM

bsakaguchi referenced this issue

2026-10-29 21:30:49 +00:00

[PR #819] [MERGED] chore: bump version to 0.3.6 in pyproject.toml #1845