[STF] Add C and Python execution-grid reshape bindings - #9899
[STF] Add C and Python execution-grid reshape bindings#9899caugonnet wants to merge 705 commits into
Conversation
Clarify how the FHE smoke tests demonstrate explicit task dependencies and the decorator-based integration path without changing test behavior.
Keep stream-context construction compatible with active CUDA graph capture and turn the back-to-back stream-context repro into an asserting regression test.
Replace direct libcudart.so ctypes loads with cuda.bindings.runtime wrappers for graph DOT export, memcpy, and stream synchronization.
The python_experimental (STF) test job for py3.13 needs the base cuda.cccl wheel for the same py_version as a runtime dependency, but the pull_request matrix only built that wheel for py3.10 and py3.14 after NVIDIA#9 reduced the base python coverage. The result was that both "Test cuda.stf (CTK12.9/CTK13.2)" jobs failed after ~4m trying to download a non-existent wheel-cccl-linux-amd64-py3.13 artifact. Add a single matrix entry that schedules the base python test (and thereby produces the wheel) at py3.13 on Linux/gcc13, mirroring the footprint of the python_experimental entry directly below it.
Fold the standalone cuda-cccl-experimental Python package back into cuda-cccl and expose CUDASTF as cuda.stf._experimental, mirroring the existing cuda.coop._experimental precedent. This removes the second wheel build and CI matrix that had no real maintenance benefit. Highlights: - Move STF module files to python/cuda_cccl/cuda/stf/_experimental/ and add an empty cuda/stf/__init__.py namespace package (no re-exports). - Simplify the STF bindings shim to match cuda.compute._bindings (drop the _BINDINGS_AVAILABLE fallback; raise ImportError on missing extension). - Build cccl.c.experimental.stf and the _stf_bindings_impl Cython extension from python/cuda_cccl/CMakeLists.txt, installing them under cuda/stf/_experimental/cuXX[/cccl]. STF is gated off on Windows. - Extend merge_cuda_wheels.py to merge both cuda/compute/cuXX and cuda/stf/_experimental/cuXX directories. - Move STF tests/probes/benches to python/cuda_cccl/tests/stf/ and rewrite imports to cuda.stf._experimental. - Rewire CI: drop build_py_experimental_wheel, point test_py_stf at build_py_wheel, rename the test script to ci/test_cuda_stf_python.sh, and collapse python_experimental to depend on the python project in ci/project_files_and_dependencies.yaml. - Update docs/python/stf.rst, docs/python/index.rst, and AGENTS.md to the new import path and install instructions (pip install cuda-cccl[cu13]). - Delete the python/cuda_cccl_experimental/ tree and associated CI scripts. This is a non-backwards-compatible migration: users on `pip install cuda-cccl-experimental[cu1X]` and `import cuda.stf as stf` must move to `pip install cuda-cccl[cu1X]` and `import cuda.stf._experimental as stf`.
The probe_*.py and bench_*.py files under python/cuda_cccl/tests/stf/ were
workshop fragments used to chase specific bugs and run one-off perf sweeps.
They use `def run(...)` rather than `test_*` functions, so pytest does not
collect them, and no maintained test or example imports them.
Delete the 22 probe_*.py files, the 4 bench_*.py files, and the
burger_scaling_smoke.{json,png} smoke artifacts. Also drop dead references
to those files from llm_helpers.py, test_llm_decode_loop.py, and
test_node_stf.py docstrings/comments.
The shared helpers (llm_helpers, pytorch_task, numba_helpers, numba_task,
numba_decorator) and the example_*.py files referenced from the docs are
kept untouched.
`branch_while_cuda_graph.dot` and `branch_while_cuda_graph.pdf` were checked in by mistake alongside the example `example_stackable_branch_while_warp.py`. The .dot was a debug rendering of one specific run; the .pdf is the compiled output. Neither is referenced by code, tests, or docs.
`python_experimental` was the matrix-project alias for the previous
standalone `cuda-cccl-experimental` wheel. Once STF moved under
`python/cuda_cccl/cuda/stf/_experimental/`, the alias had no source
files of its own and existed only to schedule `test_py_stf`.
This commit folds the STF tests directly into the `python` project:
* Removed the `python_experimental` block from
`ci/project_files_and_dependencies.yaml` (fixes the
`inspect_changes.py` SystemExit caused by an empty `include_regexes`).
* In `ci/matrix.yaml`, every `{project: 'python_experimental', jobs: ['test']}`
row becomes `{project: 'python', jobs: ['test_py_stf']}`, the project
alias under `projects:` is gone, and the MSVC exclusion is rewritten
to target `jobs: ['test_py_stf']`.
* Regenerated `ci/test/inspect_changes/*.output` (all 10 fixtures pass).
The one remaining diff from `origin/main` is `c2h_dependency.output`
picking up `python` in LITE_BUILD, which is intentional now that
`cccl_c_stf` is a lite dependency of the `python` project.
Drop exploratory Python STF demos that are not suitable as CCCL unit tests, while keeping focused API coverage and the ODE integration example.
Drop a local two-step Warp/STF experiment that is not suitable as CCCL unit-test coverage.
Splits python/cuda_cccl/tests/stf into three buckets and promotes the
Numba/PyTorch task adapters into a real package surface so they stop
living in tests/.
* New package: cuda.stf._experimental.interop.{numba, pytorch}, with
lazy optional imports. Importing cuda.stf._experimental no longer
pulls in Numba, PyTorch, or pytest.
* tests/stf/ root: pure STF unit/API tests. Renames drop the now-
redundant "stackable" qualifier (graph_scope / launchable_graph /
nested_scopes).
* tests/stf/interop/: STF<->external-runtime tests (Numba/PyTorch/Warp
adapters, @jit decorator, cuda.compute integration, capture/legacy
adoption). FDTD collapses to a single test_fdtd.py.
* tests/stf/examples/: runnable demos with test_/example_ prefixes
dropped and a main() entry point. burger_reference.py kept as the
one intentionally non-STF baseline. tests/test_examples.py extended
to discover ("STF", "stf/examples").
* Drops example_tiled_edit_distance.py, test_burger_stackable_fast.py,
and the two redundant FDTD variants.
- test_from_context: importorskip("cupy") before device/green-context
setup so a missing cupy skips without allocating a green context.
- test_numba: drop dead return after pytest.skip().
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- cg.py: drop the host_launch docstring bullet, the example never calls it - potri.py: bracket only the POTRI submission with device synchronizes and time.perf_counter(), matching cholesky.py, so verification and finalize no longer inflate the reported GFLOPS - _stf_bindings_impl.pyx: use the STF_WHILE_COND_MAX_TERMS C constant instead of hardcoded 8s for while-condition term limits Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
If _get_stream_pointer raised on a malformed stream argument after the _PrimaryContextPin was acquired, context.__cinit__ aborted with _ctx still NULL, and since _PrimaryContextPin.__dealloc__ is a deliberate no-op the primary-context retains leaked until process exit. Resolve the stream pointer first so argument validation cannot strand the pin. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 846c0b46-c1a1-452a-a0f8-afd66f03bc16
📒 Files selected for processing (10)
c/experimental/stf/include/cccl/c/experimental/stf/stf.hc/experimental/stf/src/stf.cuc/experimental/stf/test/test_placement.cppdocs/python/stf.rstpython/cuda_stf/cuda/stf/_experimental/__init__.pypython/cuda_stf/cuda/stf/_experimental/_stf_bindings.pypython/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyxpython/cuda_stf/cuda/stf/_experimental/device_array.pypython/cuda_stf/tests/stf/test_composite_places.pypython/cuda_stf/tests/stf/test_placement.py
🚧 Files skipped from review as they are similar to previous changes (4)
- python/cuda_stf/cuda/stf/_experimental/_stf_bindings.py
- python/cuda_stf/cuda/stf/_experimental/init.py
- python/cuda_stf/tests/stf/test_composite_places.py
- python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx
| tiles = (16, 16) # tile grid, distributed | ||
| tile = (512, 1024) # per-tile payload, 2 MiB of float32: page-exact | ||
| shape = tiles + tile | ||
|
|
||
| grid = stf.exec_place_grid.create(places, grid_dims=(2, 2)) | ||
|
|
||
| partition = stf.cute_partition.from_spec( | ||
| shape, | ||
| (("blocked", 0), ("blocked", 1), None, None), |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
important: Define places and make its cardinality match (2, 2). Line 325 uses an undefined variable; Line 282 only creates a two-place grid, while a 2×2 grid requires four place entries. As per path instructions, documentation examples must be technically accurate and buildable.
Source: Path instructions
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The Cython extension calls cudaStreamSynchronize directly (nogil CAI producer-stream sync) but links only the CUDA driver, so it relies on libcudart already being present in the process. Preload it through cuda.pathfinder alongside nvrtc and nvJitLink so a bare import works. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The Cython extension calls cudaStreamSynchronize directly (nogil CAI producer-stream sync) but links only the CUDA driver, so it relies on libcudart already being present in the process. Preload it through cuda.pathfinder alongside nvrtc and nvJitLink so a bare import works. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ons-python # Conflicts: # python/cuda_stf/cuda/stf/_experimental/_stf_bindings.py
cudaStreamSynchronize in _sync_cai_producer_stream was the extension's single CUDA runtime call, yet the module links only the CUDA driver, so imports relied on another package having already loaded cudart into the process. Call cuda.bindings.runtime instead (registration is not a hot path) and drop the RTLD_GLOBAL cudart preload workaround: the compiled bindings' symbol table is now honestly driver-only. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The assert guards a host>=13.4 / NVRTC<13.4 divergence (lookahead scan on sm_120). Our stacks build and run with matching 13.4 host/NVRTC on pre-sm_120 devices, where the existing lookback forcing is sufficient, so unblock the cuda_cccl build on CTK 13.4 for the localization work. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
c/experimental/stf/src/stf.cu (1)
810-819: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winsuggestion:
spec[d].policyis cast straight todim_policywith no range check, so a caller passing a garbage int from the FFI boundary produces an out-of-range enum value; any subsequentswitchon it is UB rather than a cleanNULLreturn. Validate against the known enumerators before the cast (and consider rejectingrank > 4explicitly rather than relying onmake_partition_descriptorthrowing).
♻️ Duplicate comments (1)
docs/python/stf.rst (1)
369-388: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: The tiled snippet still references undefined names:
placesat Line 369 (and its cardinality must be 4 to matchgrid_dims=(2, 2)), plusspecat Line 387, which is never bound — the specification is written inline at Line 373. Bind both so the example is copy-runnable.+ places = [stf.exec_place.device(d) for d in range(4)] grid = stf.exec_place_grid.create(places, grid_dims=(2, 2)) + spec = (("blocked", 0), ("blocked", 1)) partition = stf.cute_partition.from_spec( shape, - (("blocked", 0), ("blocked", 1), None, None), + spec + (None, None), grid.dims, )As per path instructions, documentation examples must be buildable and technically accurate.
Source: Path instructions
🧹 Nitpick comments (1)
c/experimental/stf/src/stf.cu (1)
211-212: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion:
stf_fill_placement_outputsis an internal helper but has external linkage and a public-lookingstf_name, so it leaks into the shared library's symbol table alongside the real C API. Give it internal linkage (or move it into the anonymous namespace used for the other helpers).-int stf_fill_placement_outputs( - const localized_stats& stats, const exec_place& grid, stf_placement_stats* out_stats, uint64_t* bytes_per_grid_index) +namespace +{ +[[nodiscard]] int fill_placement_outputs( + const localized_stats& stats, const exec_place& grid, stf_placement_stats* out_stats, uint64_t* bytes_per_grid_index)(and close the namespace after the function, updating the two call sites at Lines 751 and 787).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: dcd5e71f-a6b1-4d23-b796-189e1f70e486
📒 Files selected for processing (14)
.github/CODEOWNERSAGENTS.mdc/experimental/stf/include/cccl/c/experimental/stf/stf.hc/experimental/stf/src/stf.cuc/experimental/stf/test/test_placement.cppc/parallel/src/scan.cuci/matrix.yamlci/project_files_and_dependencies.yamlci/util/workflow/get_wheel_artifact_name.shdocs/python/api_reference.rstdocs/python/index.rstdocs/python/setup.rstdocs/python/stf.rstdocs/python/stf_api.rst
💤 Files with no reviewable changes (2)
- docs/python/api_reference.rst
- docs/python/index.rst
🚧 Files skipped from review as they are similar to previous changes (8)
- ci/project_files_and_dependencies.yaml
- ci/util/workflow/get_wheel_artifact_name.sh
- .github/CODEOWNERS
- AGENTS.md
- docs/python/setup.rst
- c/experimental/stf/test/test_placement.cpp
- docs/python/stf_api.rst
- ci/matrix.yaml
| int stf_cute_partition_owner(stf_cute_partition_handle h, const stf_pos4* data_coords, stf_pos4* out_grid_pos) | ||
| { | ||
| _CCCL_ASSERT(h != nullptr, "partition handle must not be null"); | ||
| _CCCL_ASSERT(data_coords != nullptr, "data_coords must not be null"); | ||
| _CCCL_ASSERT(out_grid_pos != nullptr, "out_grid_pos must not be null"); | ||
| try | ||
| { | ||
| const pos4 coords(data_coords->x, data_coords->y, data_coords->z, data_coords->t); | ||
| const pos4 owner = from_opaque_const(h)->owner(coords); | ||
| out_grid_pos->x = owner.x; | ||
| out_grid_pos->y = owner.y; | ||
| out_grid_pos->z = owner.z; | ||
| out_grid_pos->t = owner.t; | ||
| return 0; | ||
| } | ||
| catch (const ::std::exception& e) | ||
| { | ||
| fprintf(stderr, "stf_cute_partition_owner failed: %s\n", e.what()); | ||
| return 1; | ||
| } | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
important: Missing catch (...). Every other entry point in this file (stf_data_place_allocate, stf_placement_evaluate, stf_placement_evaluate_partition, ...) pairs the std::exception handler with a catch-all; here a non-std::exception throw propagates out of an extern "C" function and terminates instead of returning the documented nonzero.
catch (const ::std::exception& e)
{
fprintf(stderr, "stf_cute_partition_owner failed: %s\n", e.what());
return 1;
}
+ catch (...)
+ {
+ fprintf(stderr, "stf_cute_partition_owner failed: unknown exception\n");
+ return 1;
+ }
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| int stf_cute_partition_owner(stf_cute_partition_handle h, const stf_pos4* data_coords, stf_pos4* out_grid_pos) | |
| { | |
| _CCCL_ASSERT(h != nullptr, "partition handle must not be null"); | |
| _CCCL_ASSERT(data_coords != nullptr, "data_coords must not be null"); | |
| _CCCL_ASSERT(out_grid_pos != nullptr, "out_grid_pos must not be null"); | |
| try | |
| { | |
| const pos4 coords(data_coords->x, data_coords->y, data_coords->z, data_coords->t); | |
| const pos4 owner = from_opaque_const(h)->owner(coords); | |
| out_grid_pos->x = owner.x; | |
| out_grid_pos->y = owner.y; | |
| out_grid_pos->z = owner.z; | |
| out_grid_pos->t = owner.t; | |
| return 0; | |
| } | |
| catch (const ::std::exception& e) | |
| { | |
| fprintf(stderr, "stf_cute_partition_owner failed: %s\n", e.what()); | |
| return 1; | |
| } | |
| } | |
| int stf_cute_partition_owner(stf_cute_partition_handle h, const stf_pos4* data_coords, stf_pos4* out_grid_pos) | |
| { | |
| _CCCL_ASSERT(h != nullptr, "partition handle must not be null"); | |
| _CCCL_ASSERT(data_coords != nullptr, "data_coords must not be null"); | |
| _CCCL_ASSERT(out_grid_pos != nullptr, "out_grid_pos must not be null"); | |
| try | |
| { | |
| const pos4 coords(data_coords->x, data_coords->y, data_coords->z, data_coords->t); | |
| const pos4 owner = from_opaque_const(h)->owner(coords); | |
| out_grid_pos->x = owner.x; | |
| out_grid_pos->y = owner.y; | |
| out_grid_pos->z = owner.z; | |
| out_grid_pos->t = owner.t; | |
| return 0; | |
| } | |
| catch (const ::std::exception& e) | |
| { | |
| fprintf(stderr, "stf_cute_partition_owner failed: %s\n", e.what()); | |
| return 1; | |
| } | |
| catch (...) | |
| { | |
| fprintf(stderr, "stf_cute_partition_owner failed: unknown exception\n"); | |
| return 1; | |
| } | |
| } |
| ``stf.cond(lresidual, ">", tol)``, which also suits generated code; both | ||
| forms lower onto a single condition task and one tiny kernel regardless of | ||
| the number of terms (at most 8, one ``&``-chain or one ``|``-chain per | ||
| condition). Use ``and`` / ``or`` and these expressions raise ``TypeError``. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
suggestion: Sentence is ungrammatical and reads as if and/or are recommended.
-condition). Use ``and`` / ``or`` and these expressions raise ``TypeError``.
+condition). Using ``and`` / ``or`` on these expressions raises ``TypeError``.As per path instructions, documentation changes must be technically accurate.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| condition). Use ``and`` / ``or`` and these expressions raise ``TypeError``. | |
| condition). Using ``and`` / ``or`` on these expressions raises ``TypeError``. |
Source: Path instructions
The layer both consumer prototypes predicted and vendored (vllm localization/phase0/localized_torch.py, pytorch torch/cuda/_localized/ _alloc.py — same lineage): torch tensors backed by composite VMM data places, one ordinary contiguous tensor with pages striped by a partition. Two tiers: structured cute-partition SPEC (preferred: hashable, self-extents, drives placement_evaluate) and the Python-mapper callback escape hatch (flat byte domain, for assignments the grammar cannot express). Metadata keyed by base storage pointer so it survives views/reshape/nn.Parameter — the enabler for compiler-side detection. bf16/fp8 via same-size storage dtypes, viewed back. Lifetime is registry-pinned with explicit release(): writing the tests exposed that the consumer prototypes' weakref.finalize(buf, ...) was UNREACHABLE (the registry kept the buffer alive through the meta), so a gc-tied lifetime was never real; pinning matches the create_weights / process-lifetime weights use case and release() evicts explicitly. 6 tests: structured round-trips (f32/bf16), meta survival across views/Parameter, registry pin+release, structured-vs-callback placement parity (blocked, accuracy 1.0 on page-aligned rows), parameter wrapper, spec/mapper exclusivity. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DLPack is the ownership-carrying companion to the CUDA Array Interface: __dlpack__/__dlpack_device__ export a "dltensor" capsule whose deleter (C layer, in the Cython extension) drops a strong reference to the owning array — torch.from_dlpack(arr) yields a tensor whose storage lifetime carries the allocation, with the DeviceArray finalizer remaining the single deallocation point. CAI stays the borrowed/zero-copy path; both protocols coexist and a consumer picks by construction. Covers views (root ownership), multiple exports, unconsumed capsules, composite cute VMM places, the stream handshake (event-wait, no host block), and unrepresentable exports (structured dtypes stay CAI-only). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…nned vs DLPack-gc) localized_empty(..., lifetime="pinned"|"gc"): pinned keeps the CAI-imported, registry-pinned semantics (explicit release()); gc imports through torch.from_dlpack so the tensor's STORAGE owns the allocation -- module -> parameter -> storage -> deleter -- and the metadata registry holds no keepalive, self-evicting through a finalizer that is reachable precisely because nothing pins the buffer. localized_parameter defaults to gc: a parameter registered on a module IS the idiomatic owner, and unload paths (model swap, sleep mode) stop leaking pinned VMM. Adds the gc-tier tests (roundtrips incl. bf16 view chain, self-eviction, lifetime/placement orthogonality), a documented end-to-end example (interop/test_localized_weights_example.py), and README + stf.rst sections on the CAI/DLPack complementarity. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Brings the localized-allocation surface (interop.pytorch) and the DLPack producer on DeviceArray (CAI/DLPack dual-protocol) with lifetime tiers. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Summary
exec_place.reshape()andcollapse_axes()Python methodsextern \"C\"Stacked on #5315 and #9898. Their commits appear in this diff until they land on
main; the bindings-specific change is[STF] Bind execution grid reshape operations.Test plan
pre-commit run --files c/experimental/stf/include/cccl/c/experimental/stf/stf.h c/experimental/stf/src/stf.cu c/experimental/stf/test/test_places.cpp python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx python/cuda_stf/tests/stf/test_composite_places.py docs/python/stf.rstcmake --build build/cccl-c-stf --target cccl.c.experimental.stf.test.places.cpp -j 4ctest --test-dir build/cccl-c-stf --output-on-failure -R '^cccl\.c\.experimental\.stf\.test\.places\.cpp$'python -m pip install -e python/cuda_stf --no-depspython -m pytest -v python/cuda_stf/tests/stf/test_composite_places.py(24 passed)