kvcache-ai / kvcache-ai/Mooncake

[RFC]: A Device-API-based collective framework for Mooncake PG

Open
#2,716 2 comments 2 reactions 1 assignee Claimed by @UNIDY2002 View on GitHub
RFC
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.