kvcache-ai / kvcache-ai/Mooncake
[RFC]: A Device-API-based collective framework for Mooncake PG
- Dominant language
- C++
- Stars
- 6.6k
- Forks
- 1.2k
- Avg merge
- 3d 5h
- Merged PRs (30d)
- 312
Description
### Changes proposed
## Background
This work is proposed under inspiration from @caozhanhao, especially around the direction of making PG collectives more device-driven and more structurally scalable as the collective set grows.
## Motivation
The main limitation of the current Mooncake PG implementation is performance.
In the current PG path:
- the implementation cannot effectively exploit NVLink for the hot collective path we care about;
- the collective algorithms are still relatively naive compared with the structure and transport awareness expected from a high-performance TP PG;
- as a result, even when the functional path works, the performance gap against strong baselines is too large for the long-term target.
This is not only a matter of tuning one kernel. To close the gap, Mooncake PG needs a collective execution model that is more directly aligned with device transport capabilities.
## Proposal
We propose to introduce a Device API based collective framework for Mooncake PG, then land optimized collectives incrementally on top of it.
The overall direction is:
1. establish a reusable PG-side framework for device-driven collectives;
2. use that framework to land optimized collectives incrementally;
3. extend the same framework toward additional collectives and more topology-aware algorithms later.
The framework should make the following capabilities first-class:
- PG-owned device transport bootstrap and readiness
- reusable device-side routing primitives for:
- local/self writes
- same-node P2P writes
- cross-node RDMA/IBGDA publish and signal operations
- graph-safe sequence/slot infrastructure
## Why Start With A Framework PR
Although the long-term plan is broader than one PR, the first upstream step should be framework-oriented.
That first PR should focus on the shared substrate that later collective PRs will depend on, especially:
- experimental runtime/framework boundary
- bootstrap and metadata exchange flow
- device transport ownership
- graph-safe sequence-slot ownership
- reusable device-side transport/routing primitives
## Planned PR Stack
The expected shape is a stacked series of PRs rather than one large merge.
### PR1: Device API PG framework
Focus:
- experimental framework/runtime layer
- bootstrap / metadata exchange / connection sequencing
- PG-owned device transport state
- reusable device routing primitives
- graph-safe sequence-slot ownership
### PR2+: Collective implementations on top of the framework
Focus:
- collective-specific host orchestration
- collective-specific device kernel paths
- correctness and performance validation
Later PRs can extend the same structure toward:
- improved all-reduce paths
- hierarchical or topology-aware algorithms
- broader collective coverage
- further transport-aware performance tuning
cc: @yuechen-sys
### Before submitting a new issue...
- [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)
Contributor guide
Assessment
This issue has not been assessed yet.