[RFC] An Extensible Collective Framework for Mooncake PG

## Summary

This RFC proposes an extensible GPU collective framework for Mooncake PG. The framework separates PG-level backend dispatch from group-wide algorithm selection, rank-local planning, GPU execution, and device data movement, while preserving the existing PG APIs and elastic semantics.

This proposal builds on #2716 by turning its device-collective direction into a concrete framework for GPU collectives. The framework is introduced incrementally alongside the legacy implementation, and the new collective path remains opt-in during the rollout.

## Motivation

The legacy GPU collective path couples collective protocols with host workers, transfer submission, synchronization, and failure handling. Adding several algorithms per operation and tuning them for different message sizes and topologies would otherwise duplicate these concerns and make cross-rank consistency harder to maintain.

The new framework gives GPU collective algorithms a narrower responsibility: implement a protocol against a rank-local plan and a device transfer interface. Membership, capability negotiation, backend dispatch, algorithm selection, resource ownership, and recovery remain framework responsibilities.

## Goals

- Define reusable planning and execution boundaries for GPU collective operations.
- Define algorithm selection within the new path so every rank receives the same policy and runs a consistent collective protocol.
- Decouple device collective algorithms from route-specific transfer details.
- Allow membership or algorithm changes to reuse a CUDA Graph capture.
- Migrate incrementally with explicit opt-in while the legacy path remains the default.

## Non-Goals

- Replacing all legacy collectives during the initial rollout.
- Switching to another route after a transfer failure in the initial implementation.

## Architecture

```
+-------------------+
| torch.distributed |
+---------+---------+
|
v
+===================================================================================================+
| Mooncake PG |
| |
| +-------------------------+ |
| | Mooncake Communicator | |
| +------------+------------+ |
| | |
| +--------------------------------+-----------------------+ |
| | | | |
| v v v |
| +-------------------+ +---------------------------------------------------+ |
| | Legacy Collective | | New Collective Runtime | |
| | Runtime | | | |
| +---------+---------+ | +----------------+ +--------------+ | |
| | | | Device Runtime | | CPU Runtime |------+-------+ |
| | | +--------+-------+ +--------------+ | | |
| | +----------|----------------------------------------+ | |
| | | | |
| | v | |
| | +-----------------------------------------------------------+ | |
| | | Device Transfer Service | | |
| | | | | |
| | | +-----------+ +------------+ +-----------+ +------------+ | | |
| | | | P2P Route | | RDMA Route | | GIN Route | | Host Proxy | | | |
| | | +-----+-----+ +-----+------+ +-----+-----+ +------+-----+ | | |
| | +-------|-------------|--------------|--------------|-------+ | |
| | | | | | | |
+=============|================|=============|==============|==============|=================|======+
| +-------------+--------------+ | |
| | | |
| v v |
| +--------------------------+ +------------------------+ |
| | Transfer Engine (Device) | | Transfer Engine (Host) |<---+
| +--------------------------+ +------------------------+
| ^
| |
+------------------------------------------------------------+
```

The public collective APIs remain unchanged. The Coordinator owns authoritative membership and publishes group-wide decisions. `MooncakeCommunicator` dispatches operations to the selected legacy or new backend. Within the new backend, a separate algorithm policy chooses a compatible algorithm, the GPU runtime materializes its rank-local execution state, and the kernel moves data through `DeviceTransferService`.

### Walkthrough

The diagram separates API dispatch, collective execution, and data movement. It shows the intended component boundaries rather than the scope of the first PR.

| Component | Role in the diagram | Main connection |
| --- | --- | --- |
| `MooncakeCommunicator` | Keeps the public PG API unchanged and dispatches each collective to the group-wide selected backend. | Sends an operation to either the legacy runtime or the new collective runtime. |
| Legacy Collective Runtime | Represents the existing implementation that remains available during migration. | Uses the host-facing Transfer Engine path. |
| New Collective Runtime | Groups new collective implementations behind a boundary separate from the legacy path. | Contains the GPU-oriented Device Runtime and leaves room for a separate, future CPU Runtime. |
| Device Runtime | Runs GPU collective algorithms and owns the GPU-side execution state needed by those algorithms. | Launches GPU work and accesses data movement only through `DeviceTransferService`. |
| CPU Runtime | Marks possible future CPU collective work. | Uses a host-oriented execution and transfer path rather than `DeviceTransferService`. |
| Device Transfer Service | Hides route-specific transport details from GPU collective algorithms and presents one device-facing data movement interface. | Resolves a usable route for each peer and connects the Device Runtime to route providers. |
| P2P, RDMA, and NCCL GIN routes | Represent device-initiated ways to reach a peer. Their availability depends on hardware. | Implement the common transfer-service contract using their route-specific mechanisms. |
| Host Proxy route | Bridges a device-initiated operation to transport work that must be submitted by the host. | Connects `DeviceTransferService` to the host-facing Transfer Engine path. |
| Transfer Engine (Device/Host) | Represents the underlying device- and host-facing data-movement facilities used by both the new and legacy code. | Performs the actual transfer underhood. |

The control plane is omitted from the diagram. It publishes membership and the backend decision consumed by `MooncakeCommunicator`. A separately published algorithm policy is used only after dispatch reaches the new Device Runtime, as described below. The main data path is therefore `MooncakeCommunicator` → collective runtime → transfer service → transfer engine.

## Backend Selection

Backend selection is a PG-level dispatch decision between the legacy collective path and the new collective path. It is primarily a rollout and compatibility boundary; it does not select a concrete algorithm within the new backend.

The decision must be group-wide because the two backends may run different protocols. Each GPU rank reports one backend preference together with its capability for the new backend. The preference is read when a GPU communicator is created from `MOONCAKE_PG_PREFERRED_GPU_COLLECTIVE_BACKEND`: an unset value or `legacy` selects the default legacy preference, while `new` explicitly requests the new path. Any other value is a configuration error.

The environment variable expresses only a local preference. The Coordinator selects the new backend only when every active rank requests it and has the required capability; any preference mismatch or capability gap selects the legacy backend for the whole group. The result is published in the `GroupView` at group bootstrap and after membership transitions. All worker processes must therefore set the variable consistently before creating their communicators, and changing it does not affect an existing communicator.

## Algorithm Selection

Algorithm selection happens after the new backend has been selected. It chooses a concrete implementation for a collective invocation, such as different AllReduce algorithms for different message-size ranges or topologies.

Algorithm selection must be group-wide because ranks running different collective protocols can wait for different peers or signals and deadlock or corrupt data. The Coordinator already owns authoritative membership and receives rank capabilities, so it is the natural place to resolve a common policy at group bootstrap and membership transitions.

The policy maps collective and message-size ranges to algorithms, using inputs such as operation, active membership, topology, available routes, and device resource constraints. The Coordinator publishes the same policy to every rank, and invocation metadata indexes it locally without a control-plane RPC.

## Device Transfer Service

`DeviceTransferService` owns a stable registered device region, a device-visible route table, fixed transfer lanes, and route-provider resources. A device arena suballocates shared workspace and communicator control slices from that region, making their addresses and memory cost predictable.

The device-facing contract provides payload-plus-signal, standalone signal, signal wait, and completion semantics. Algorithms operate on peers and registered-region offsets; route providers handle peer mappings, transport metadata, queueing, and host/device completion differences.

Route selection is separate from both backend and algorithm selection. Route providers publish their capabilities, and the runtime materializes one reachable route per peer in the device-visible route table. P2P performs direct GPU-accessible transfers; host proxy lets the GPU enqueue commands that a host worker submits through the TE host API; additional route providers can implement the same transfer contract.

## Failure Recovery

Device code can detect lack of progress but cannot run PG's host-side membership protocol. A process-wide `DeviceCollectiveRecoveryWorker` bridges the two through a generation-based device/host mailbox:

1. A timed-out transfer records the failed in-group rank, and the invocation converges at a device safe point where no channel can continue using the old plan or shared transfer buffers.
2. After that safe point, one CTA publishes a failure generation and waits. The failed invocation does not read a replacement plan and therefore cannot race with host-side rematerialization.
3. The recovery worker drains outstanding host-proxy work, converts the failed rank into the existing PG link-failure event, and enters `syncAfterFailure`.
4. The Coordinator reconciles failure reports and publishes an authoritative `GroupView`. The Agent applies it, and the selected device runtime rematerializes its plan for the new active membership.
5. Only after recovery state has been published does the worker acknowledge the generation. The failed invocation exits, and the next eager launch or a graph replay loads the current plan.

### CUDA Graph Replay After Membership Changes

Captured kernels can reference stable runtime objects such as the plan slot, transfer-service handle, route table, protocol storage, recovery mailbox, and workspace. Keeping membership-dependent information in the contents of those objects makes it possible to reuse a capture after membership changes.

## CPU Collectives

CPU collectives remain planned future work, but they are not another backend of the GPU framework described above. A CPU design may reuse mechanisms such as `GroupView`, Coordinator-driven selection, and communicator dispatch, while defining its own cpu-oriented algorithm, execution, transport, and recovery abstractions.

## Rollout Plan

The new collective path is opt-in throughout the initial rollout. The legacy path remains the default; changing that default requires an explicit decision after correctness, compatibility, and performance criteria have been met.

| Stage | Scope | Status |
| --- | --- | --- |
| Phase 1 | Establish the framework under `device_comm`, opt-in group-wide backend dispatch, one CUDA AllReduce implementation, P2P and host-proxy routes, recovery integration, and coexistence with the default legacy implementation. | #3484 |
| Phase 2 | Introduce group-wide algorithm selection, add more AllReduce algorithms and an RDMA route, derive and tune selection thresholds, and optimize AllReduce latency and bandwidth. | Planned |
| Phase 3 | Add more collective operations, such as AllGather and ReduceScatter. | Planned |
| Phase 4 | Add and tune multiple algorithms for common collectives, including message-size- and topology-aware selection. | Planned |
| Phase 5+ | Design and introduce CPU collectives as a host-oriented path, reusing PG-level control-plane pieces where useful, and migrate CPU operations incrementally. | Future |

Start by reading the existing MooncakeCommunicator and Coordinator paths, then trace the legacy collective runtime and the DeviceTransferService boundary described here. The RFC does not identify implementation files, tests, or a first-PR scope, so completion criteria must be agreed before work can begin.

Written by the indexing model from the issue text.