[sharded 7/7] CUDA graph capture: the contract, enforced and tested

[caugonnet]

Stacked on #10 (sharded/sort-va); review only the last commit here.

This PR establishes and tests the CUDA-graph-capture contract of the sharded surface, and adds the small enforcement it needs. A new CUDA graph capture section in sharded.rst is the user-facing statement of the contract.

The contract

What captures. The non-blocking elementwise family (fill, sequence, iota, tabulate, generate, for_each, transform with blocking = false) is pure per-shard kernel launches on the shard streams, and captures with the standard fork/join event idiom (places::make_stream_wait_for in both directions around the recorded pipeline). Properties pinned by tests:

• the per-place SM confinement of the (green-context) shard streams survives inside the instantiated graph: kernels captured from two locality-domain streams land on disjoint SM sets on replay (SM-id check);
• replays recompute from the current shard contents, so inputs can be rewritten between launches;
• cross-stream dependencies recorded inside the capture (out-of-place transforms across stream colors) work;
• contiguous (allocate_contiguous) arrays are transparent to capture — the VMM mappings pre-exist the graph — and per-shard stages compose with whole-array kernels through contiguous_data() in one graph.

What refuses, and how. Everything that allocates containers, transfers host data or synchronizes cannot be represented in a graph (the synchronous algorithms all stage per-place partials through the host). These now throw std::runtime_error when invoked under an active capture, detected by a safe query before any CUDA call that would invalidate the capture (or abort through cuda_safe_call): container allocate/allocate_contiguous, copy_from_host/copy_to_host/copy_between, sharded_array::sync, place_group::sync/sync_stream, and reduce/sum/min/max, the scans, count/count_if, histogram_even, copy_if/filter/remove_if, unique, adjacent_difference, and sort (either engine). The refused capture stays valid and keeps accepting supported work; everything refused works eagerly after the capture ends. The new helper places::stream_in_capture(stream) also probes the legacy stream, so the guards fire when a global-mode capture is active anywhere in the process (where these operations would invalidate it under the CUDA capture rules).

Benign by construction (also tested): shard adoption and slice (host-only bookkeeping, usable during capture), and place_group construction / lazy stream materialization (records nothing into the graph).

Graph-owned memory. place_memory_resource stream-ordered allocation is deliberately left capturable: an allocate on a capturing stream records a graph memory node drawing from the place's pool. Tests pin both shapes — a balanced alloc/free pair enclosed in the capture instantiates and replays freely; an allocation not freed in-graph stays live after a launch (the captured pointer reads back), an immediate relaunch fails predictably, and freeing the pointer outside the graph with cudaFreeAsync re-arms the launch. The documented user contract stays simple: allocate outside capture; capture computation only.

Why the guards throw instead of letting the driver error surface

Under a global-mode capture, the first illegal call (cudaStreamSynchronize, pinned-host allocation, ...) both returns cudaErrorStreamCaptureUnsupported and invalidates the capture — and several of these paths run under cuda_safe_call, which aborts. The entry guards convert that into a clean, catchable refusal that leaves the capture usable.

Replay smoke (2 locality domains, GB300, CUDA 13.4)

A 6-kernel captured pipeline replayed 1000x: ~12.4 us/iteration for the eager dispatch loop vs ~6.2 us/iteration for graph replay in the launch-bound regime (min of 3, stable across runs) — about a 2x cut in per-iteration launch overhead, the figure that matters for host-dispatch-bound loops.

Tests

cudax/test/sharded/graph_capture/{elementwise_pipeline, reduce_scan_capture, must_not_capture, graph_owned_memory, contiguous_capture}.cu — silent on success, following the suite style. 36/36 places+sharded tests green on 2 locality domains (sm_103a, -DCMAKE_CUDA_ARCHITECTURES=103a, NCCL off).

---

Generated with Claude Code