cudax sharded: places communicator -- a places-backed model of the multi-GPU communicator concept + the engine-bindings adapter

[caugonnet]

Description

Motivation

The MGMN algorithms (__multi_gpu) program the ranks rung against a
duck-typed communicator concept, with NCCL as the shipped model. Code
written against that concept cannot currently run on the rung below —
several places inside one process — even though nothing in the concept
requires a network: a second, in-process model would let users prototype
rank-scoped code on one node (e.g. two locality domains) and scale it
unchanged to multi-process NCCL ranks, and would let in-process
containers drive the MGMN constructs as engines.

This PR connects the two rungs at the top of the cooperation-scope
ladder. The MGMN algorithms (__multi_gpu) program the ranks rung
against a duck-typed communicator concept; this PR adds a
places-backed model of that concept, so the same constructs run
unmodified over in-process places (locality domains, green-context
places, devices-as-places) — and the adapter that lets sharded
containers drive them.

What's in the PR

• __sharded/communicator.cuh:
  • places_communicator / basic_places_communicator: each place of
    a place_group is one rank (rank = place index, native_handle()
    = the place). Because the places of one process share a virtual
    address space, the verbs lower to device-to-device copies
    (send/recv rendezvous-matched inside a group guard,
    all_gather[_v], all_to_all[_v]) and all_reduce to a single
    fold kernel that combines every rank's partial in fixed rank
    order — bit-identical results run to run for a fixed place list.
    The basic variant omits all_reduce, so both MGMN combine paths
    (direct all_reduce, and all_gather + local combine) stay reachable
    and tested.
  • logical_device() returns a place-backed model of the
    {context(), underlying_device()} pair the MGMN algorithms
    activate and select resources through.
  • make_communicators(place_group&) (and an explicit-place-list
    overload): one communicator per place, one shared group state.
    Groups spanning devices are refused with a diagnostic — that
    composition belongs to a hierarchical transport tier, not to this
    rung.
  • Concept conformance stated in code: static_asserts against
    the __multi_gpu concepts (__communicator, __has_all_reduce
    present/absent per variant, __has_all_gather[_v],
    __has_all_to_all[_v]).
  • bind_engine(group, sharded_array): the two-tier seam. The
    container tier owns placement and resources; the engine tier owns
    cross-place choreography. This adapter manufactures what an MGMN
    engine consumes — one communicator, environment (stream + per-place
    memory resource), iterator and size per shard.
• __places/place_group.cuh: place_memory_resource gains
  default_queries (device_accessible) and host/device-annotated
  stream-ordered allocate/deallocate (terminating on device, the CUB
  device_memory_resource idiom), so the per-place resources travel
  through environments into CUB_RUNTIME_FUNCTION dispatch and
  cuda::buffer-based engine temporaries. Net effect: engine scratch
  is allocated on the place whose rank uses it.
• Only the vendor-free __multi_gpu/concepts.h is included; the
  umbrella sharded.cuh stays buildable with NCCL disabled.

Tests (cudax/test/sharded/communicator/)

• concepts.cu: the conformance static_asserts + factory contract
  (rank/size/order, native handle identity, empty-list refusal).
• verbs.cu (2 locality domains): send/recv rendezvous, all_gather
  correctness, all_reduce correctness on random fp32 + broadcast
  bitwise agreement across ranks + bit-determinism over 5 repeated
  runs (the fixed fold order, checked).
• bridge.cu — the conformance evidence as CI:
  cuda::experimental::reduce (the MGMN range API) running over
  communicators and environments manufactured from a sharded_array
  by bind_engine, on 2 locality domains with uneven shards, checked
  against a single-place CUB DeviceReduce reference and an fp64 host
  reference, through both combine paths, with bit-determinism
  across 5 runs per path.

Verified on a 2-locality-domain GB300 node (sm_103a, CUDA 13.4), built
with -Dcudax_ENABLE_NCCL=OFF.

Checklist

• New or existing tests cover these changes.
• The documentation is up to date with these changes.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.