Use GPU AllreduceV for distributed quantile sketch exchange

[RAMitchell]

Summary

This changes the distributed GPU quantile sketch exchange from an allgather-style pattern to a GPU AllreduceV path built on the existing collective interfaces.

The new implementation packs the sketch into a raw byte payload, performs a variable-size reduction over NCCL, and materializes the reduced sketch back into SketchContainer. The reduction uses a tree of point-to-point edges for the reduce phase and the existing broadcast path for the final fan-out.

Motivation

The previous allgather-style exchange causes communicated bytes and temporary communication storage to grow poorly with worker count because every rank materializes peer payloads.

A tree-shaped reduction keeps the communication pattern closer to the CPU AllreduceV path and avoids that replication on the NCCL fast path.

Design

The main pieces are:

• a GPU AllreduceV API aligned with the existing collective API shape in src/collective/allreduce.h
• a private CUDA implementation in src/collective/allreduce_v.cuh
• quantile-side packing/unpacking in src/common/quantile.cu

The current NCCL implementation:

• uses existing Comm::Chan(peer) wrappers for tree Send/Recv edges
• uses the existing broadcast wrapper for the final dissemination
• preserves upstream nonblocking NCCL communicator init
• explicitly waits for communicator progress after p2p edges so the tree is stable with nonblocking NCCL communicators

Main Changes

• Add GPU AllreduceV support to the collective layer.
• Route distributed GPU quantile sketch exchange through that helper instead of the old allgather-style path.
• Keep the quantile reduction logic local to quantile.cu by packing the sketch as bytes and reducing with merge + prune semantics.
• Add focused tests for the GPU AllreduceV tree shape and repeated-call stability.

Testing

Ran:

• MGPUAllreduceVTest.SimpleTree
• MGPUAllreduceVTest.RepeatedTree
• MGPUAllreduceTest.Timeout
• GPUQuantile.Prune
• GPUQuantile.MergeBasic
• GPUQuantile.MultiMerge
• MGPUQuantileTest.AllReduceBasic

Follow-up

• Replace the temporary relaxed quantile cut-value tolerance with a check based on sketch/rank guarantees.
• Decide how AllreduceV should behave on non-NCCL distributed GPU backends such as federated CUDA communication.

[Copilot]

Pull request overview

This PR reworks multi-GPU quantile sketch exchange to use an NCCL-based tree reduction (via cached NCCL sub-communicators) instead of an allgather-style pattern, aiming to reduce communication volume and avoid per-rank replication.

Changes:

• Add cached NCCL subgroup communicator support (ncclCommSplit) in the collective communicator layer.
• Implement an NCCL tree reduction + broadcast path for SketchContainer::AllReduce that keeps sketch payloads on GPU and prunes after each merge step.
• Update MGPU quantile test validation to compare final cut pointers exactly and cut values within a relaxed tolerance.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
src/common/quantile.cu	Adds NCCL tree reduction/broadcast exchange path for GPU sketch allreduce using cached subgroups and GPU-resident payloads.
src/collective/comm.cuh	Extends NCCLComm with cached subgroup support and subgroup metadata (ranks, communicator, stream).
src/collective/comm.cu	Implements NCCLComm::GetSubgroup via ncclCommSplit, manages subgroup lifetime, and changes root communicator config to blocking.
src/collective/nccl_stub.h	Extends NCCL stub API to include ncclCommSplit.
src/collective/nccl_stub.cc	Loads/binds ncclCommSplit (dlopen and static-link paths).
tests/cpp/common/test_quantile.cu	Adjusts MGPUQuantileTest.AllReduceBasic to validate final cuts (ptrs exact, values tolerant) instead of sketch-level identity.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[trivialfis]

The NCCLComm has a set of channels stored in the Comm base class, with the usual stream block methods. Do you think we can use them?

I would like to reuse some existing wrappers, like channel and broadcast, instead of going directly to the NCCL handle, if possible.

[RAMitchell]

I've had a try and it seems to work.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[RAMitchell]

I had to include a fallback for federated learning.

[trivialfis]

Let me take a look next week.

[trivialfis]

I had to include a fallback for federated learning.

Unfortunately, the federated learning module cannot be made to fit the normal ring/tree structure, as peers are forbidden to see the messages from other peers.