python cuda_sharded: bindings for the sharded rung (bound containers, tiered algorithms, free-threading-native)

[caugonnet]

Python tier of the sharded stack (PR 1)

New package python/cuda_sharded (import cuda.sharded._experimental): Cython bindings over this branch's cuda::experimental::sharded headers, following the python/cuda_stf packaging precedent (scikit-build-core, in-tree cython, cuXX/ extension directory, cuda/ namespace packages).

Design

Bind the containers -- one implementation. place_group (by_locality_domains / by_devices, per-place streams with colors, sync) and sharded_array (allocate with total size or explicit per-shard sizes, allocate_contiguous over the VMM backing, adopt of external per-shard buffers via __cuda_array_interface__, numpy round-trips) hold opaque handles to the C++ objects; placement, the fixed-size contract and the contiguous backing live once, in the headers. Contract errors cross the boundary as exceptions (std::invalid_argument -> ValueError, std::runtime_error -> RuntimeError).

Tier the algorithms by what crosses the Python/C++ boundary.

1. Opaque algorithms (default): fill, sequence/iota, reduce (sum/min/max), inclusive_scan/exclusive_scan, adjacent_difference, count, histogram_even, sort -- one crossing per call; C++ owns the per-shard loop, the per-place streams, and the cross-place combine.
2. Operator-parameterized: transform / transform_binary with standard-op descriptors (negate, scale, add_scalar; add, mul, axpy) lowering to the pure C++ path -- the operator is an enum, not a per-element callback. Arbitrary Python operators JIT-compiled once (cuda.compute numba -> LTO-IR, cached by (op, dtype, algorithm), executed by C++ per shard) are a documented follow-up.
3. Per-shard escape hatch (always available): sharded_array.shard(i) is a zero-copy __cuda_array_interface__ view (pointer, shape, typestr, and the shard's reference stream), so numba/cupy/torch kernels run on individual shards directly.

Free-threading-native from day one: freethreading_compatible=True, no mutable module-level state, GIL released around blocking shim calls; the README carries the threading contract (shared place_group, per-thread arrays; same-array mutation unsynchronized, as in C++).

The shim is two nvcc translation units (sort isolated -- it instantiates the distributed engine and dominates compile time) explicitly instantiating f32/f64/i32/i64; the cython-generated TU is CUDA-free and compiles with the host C++ compiler. The extension links cudart statically and installs into cuXX/ like cuda-stf.

Tests (GB300, sm_103a, CUDA 13.4, 2 locality domains)

tests/sharded/: containers (round-trips for all four dtypes, uneven splits, explicit sizes, contiguous allocation checked shard-by-shard against the base pointer, shard-view CAI fields, contract ValueErrors), every tier-1 algorithm against a numpy reference (including sort over mixed distributions and on a contiguous array), transform descriptors, cupy and numba per-shard interop (external kernels write / algorithms read, and the reverse; adoption + algorithms writing through to the adopted buffers), and threaded dispatch over a shared group (8 threads x tier-1 calls; concurrent read-only calls on one array; concurrent sorts).

• Python 3.12: 80 passed, 1 skipped (the GIL-state check, N/A on a GIL build)
• Python 3.14t (free-threaded): 80 passed, 1 skipped (numba, not available on 3.14t); sys._is_gil_enabled() verified False after importing the bindings

Both ruff and clang-format/gersemi clean.

Follow-ups

• numba -> LTO-IR rung for arbitrary Python operators (tier 2b);
• sharded_csr + spmv/spmm bindings behind the cuSPARSE gate, mirroring the C++ layer;
• graph/STF capture of per-iteration sharded sequences (the loop-level tier);
• multi-CUDA wheel packaging (merge_cuda_wheels.py-style).

🤖 Generated with Claude Code