cudax sharded: a shared-address-space engine for the places-rung sort

[caugonnet]

This PR adds a shared-address-space engine for the places rung behind sharded::sort, composing with the distributed engine via the two-tier design: the container tier is unchanged, and sort_engine selects the tier-2 engine (automatic by detection, shared_va / distributed pinned explicitly for A/B or portability testing).

One name, two engines — one per rung of the ladder

CUB's invariant is that a primitive at scope N runs scope N-1's primitive locally and combines using what scope N shares. sharded::sort now instantiates that at both of the upper rungs:

rung	what it shares	engine
places (one process, one VA)	the address space: direct loads across shard boundaries	sort_engine::shared_va (this PR)
ranks (nothing shared)	message passing	sort_engine::distributed (the __multi_gpu MGMN sort via bind_engine, unchanged)

The distributed engine remains the portability path — code written against it runs unchanged from two locality domains to N NCCL ranks — and the automatic fallback wherever the shards do not share an address space.

The shared-address-space engine

When every shard lives in one device's address space (locality domains, or the device itself), the cross-place combine can be computed through loads rather than choreographed through messages:

1. Local sorts, out of place: each place sorts its shard into a per-place auxiliary run — cub::DeviceRadixSort for arithmetic keys under the default ascending/descending orders, cub::DeviceMergeSort for arbitrary comparators.
2. Exact splitters by multi-sequence selection: all sorted runs are visible in the shared address space, so the splitters at the container's fixed shard boundaries are computed exactly by one small kernel (a k-way binary search over the P runs, ties broken totally by (run, index)). No estimation loop, no tolerance: the selection is deterministic and lands exactly on the boundaries.
3. Fused gather-merge: each destination place k-way-merges the selected sub-ranges directly out of the source runs — reads across shard boundaries — and writes straight into its own shard storage. The output is at the original boundaries by construction, so contiguous (allocate_contiguous) arrays read as one sorted array through the base pointer, unchanged contract.

Synchronization is plain stream/event ordering between the three phases. Temporary footprint is one auxiliary run per place (1x the data) plus per-place merge scratch only beyond two places.

Measured (GB300, sm_103a, CUDA 13.4, 2 locality domains, fp32, min of 3)

fp32 keys, place_group::by_locality_domains() (2 places), against the whole-device cub::DeviceRadixSort baseline; min of 3 in two independent quiet-node sessions:

n	uniform: whole-device CUB	shared_va (2 domains)		lognormal: whole-device CUB	shared_va (2 domains)
4M	0.15 ms	0.20 ms		0.15 ms	0.23 ms
16M	0.43 ms	0.49 ms		0.43 ms	0.47 ms
67M	1.55 ms	1.53 ms		1.55 ms	1.49 ms
268M	6.08 ms	5.61 ms		6.10 ms	6.27 ms

At 268M uniform the two-domain engine finishes ahead of the whole-device cub::DeviceRadixSort baseline while keeping per-domain memory confinement: the local sorts run concurrently on the two domains, and radix sort scales sublinearly with SM count, so two half-size sorts beat one full-size one. Phase decomposition at 268M (steady state): local sorts ~5.1 ms, selection ~0.1 ms, gather-merge ~0.5 ms.

Three supporting fixes are part of the branch, each measured independently:

• locality-domain memory pools and (unconditionally) the device default pools now set an unlimited release threshold at creation — with the default threshold of 0, every synchronization returned pool memory to the OS and algorithm-scale scratch paid milliseconds of physical re-backing per call;
• the selection's targets travel by kernel parameter and its splits return through a pooled device buffer + one async copy (a per-call pinned host allocation dominated the phase);
• a merge-path diagonal-search fast path for the two-run case of the selection kernel.

Contract and tests

Both engines honor the same contract: each shard ends holding the slice of the globally sorted sequence at its original boundaries (sizes/offsets/capacities unchanged; the engine _CCCL_VERIFYs the exact-boundary landing), sorting is not stable, and repeated runs are bitwise identical. New test algorithms/sort_engines.cu: six distributions x sizes x float/int on BOTH engines with cross-engine byte agreement, custom comparators on both engines plus the radix descending fast path, contiguous backing on both engines, per-engine bitwise repeat-run identity, and the eligibility rules (host-backed shards refuse an explicit shared_va request). The pre-existing algorithms/sort.cu runs unchanged. Correctness was additionally verified against std::sort byte-for-byte at n = 64 .. 268M over uniform / lognormal / all-equal / pre-sorted / reverse / dup-heavy inputs, with uneven shards and the contiguous backing at 268M, plus bitwise repeat-run identity at 268M. 31/31 places+sharded ctest green; header-inclusion targets green.

🤖 Generated with Claude Code