[STF] locality domain for python - #10703
Draft
caugonnet wants to merge 706 commits into
Draft
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.
The device/green_ctx data places allocate via cudaMallocAsync, which draws from the *current* device's default pool, so they must set the device first. The locality-domain place instead calls cuMemAllocFromPoolAsync with an explicit pool created with props.location.id == devid, so placement does not depend on the current device. No caller established a matching device either: neither STF allocate call site (stream_ctx.cuh, adapters.cuh) nor the allocator layers above them set one, and the enclosing exec_place_scope from task::acquire activates the *exec* place, which may be a different device than the data place. The switch was therefore both unnecessary and unrestored, leaking the current device into caller code. Removing it also makes allocate() symmetric with deallocate(), which never switched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
- Release the retained primary context on every exit path in the fake-cache granularity probe: a throwing cuCtxGetDevResource leaked the reference. - Override is_device() to true on the fallback exec place impl: the native backend reports true through exec_place_cuda_ctx_impl, and callers that branch on is_device() (e.g. the parallel_for reduction path) must see the same answer from both backends. - Include <cstdio> (fprintf) and scope_guard.cuh (SCOPE) directly in locality_domain.cuh, and <cstdio>/<cstddef> in the tests that use fprintf/size_t, instead of relying on transitive includes. - Strengthen the exec test: cudaSetDevice(dev) beforehand made the in-scope device check pass trivially for a no-op activate(); on the native backend also verify the current driver context actually changed under the scope. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
- Make the places.rst example handle the zero-domain case with a fallback branch instead of assert, and mark it schematic (it assumes the usual ctx/lX/kernel setup from the STF introduction). - Clarify the CUDA version qualifier in stf.rst: native support is 13.4+, and the same API compiles on older toolkits with the whole-device fallback reporting a single domain. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…instead The rationale of the API is that a device without locality-domain support reports a single domain covering the whole device, so callers keep one code path and never need a zero-count special case. - The public count validates the device ordinal with EXPECT (like data_place::device) and otherwise always returns >= 1. - On a native (13.4+) build whose driver cannot answer the locality-domain query, the count degrades to 1 and places actually work: the context cache builds one green context spanning the full SM resource instead of throwing, and the memory pool / mem_create paths fall back to plain device memory (the localized location type would be rejected by such a driver) -- i.e. exactly the documented pre-13.4 fallback semantics, selected at runtime. - The fallback backend count returns 1 unconditionally (device validation moved to the public function). - Tests: zero-count waivers become EXPECT(count >= 1); invalid-device checks now expect an exception instead of 0. - Docs: drop the zero-count language and the zero-guard from the example. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
# Conflicts: # ci/matrix.yaml
… stf-locality-domain-python
Stacks on stf_c_api (NVIDIA#5315) + stf-locality-domain-places (NVIDIA#10658): the placement work shines with uGPU locality domains, so expose them where the allocation surface lives. C API: stf_locality_domain_count (never 0 for a valid device -- the whole-device degrade contract; 0 = invalid device, detail on stderr), stf_exec_place_locality_domain, stf_exec_place_locality_domain_grid (make_locality_domain_grid), stf_data_place_locality_domain. Python: stf.locality_domain_count(dev), exec_place.locality_domain(dev, d), data_place.locality_domain(dev, d), exec_place_grid.locality_domains(dev). Tests adapt to the reported count (single whole-device domain on fallback backends): count contract, place identity, per-domain allocation roundtrip, a domain-grid task reading at the replicated place (the uGPU counterpart of the replicated grid test), and per-domain pinned tasks. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…dings Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Contributor
2 tasks
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The missing entry point for "give me a grid for this machine": granularity="device" (one place per CUDA device, like from_devices over all of them) or granularity="locality_domain" (one place per locality domain of every device, device-major). A device without native support contributes its single whole-device domain, so the domain granularity degrades to the device one exactly where domains are unavailable -- one code path for callers, per the never-0 count contract. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
caugonnet
added a commit
to caugonnet/cccl
that referenced
this pull request
Aug 7, 2026
Parametrize the example suite's grid over two granularities: the repeat device grid, and the machine's locality domains (via exec_place_grid.machine(granularity="locality_domain") when the PR NVIDIA#10703 bindings are present, locality_domains(0) as fallback, clean skip otherwise). The whole programming-model spectrum -- compiled pointwise, unsplit-dim softmax, partials+fold, misalignment, CUDA-graph capture, TinyMLP -- now exercises the substrate the placement work is for. The partials assert becomes grid-size-agnostic (domain counts vary by machine). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A create()-built grid has no affine data place; a dependency without an explicit place then throws C++-side and TERMINATES through the void C API instead of raising -- the known 5315-family error-channel gap, hit here on GB300. Use the proven replicated-read pattern and document the constraint in the test. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Design decision (Cedric): blocked along dimension 0 is the natural strategy for a machine-level grid, so machine() attaches it as the grid's affine data place (native stf_partition_fn_blocked, no callback trampoline). Bare dependencies -- lX.rw() with no explicit data place -- now resolve on machine grids instead of terminating through the void C API for lack of an affine; the test exercises exactly the case that aborted on GB300, alongside an explicit replicated read. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ntyped access) Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…er test Found on GB300 while validating machine()'s default affine: a replicated read whose valid source instance lives at a composite place terminates in the device exec-place deactivate with "invalid device ordinal" -- the broadcast-copy path derives a restore ordinal from a composite affine. machine() merely makes the sequence (composite write, then replicated read) easy to reach; the interplay predates it. The reproducer is recorded in the test comment; the test sources the replicated read from the host instance until the C++ path is fixed. Bare rw at the default blocked affine works and stays covered. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
GB300 bisection: the aborts were NOT a replicated-broadcast issue (that sequence passes) but affine mutation on a SHARED place. make_grid degenerates a size-1 grid to the place itself; for the device granularity that is the process-shared exec_place::device(0), so attaching the blocked composite affine mutated the shared scalar place -- and every later deactivate derived its restore ordinal from a composite affine (cudaSetDevice(-5), invalid device ordinal). Attach the default affine only when the machine grid has more than one place; a scalar place's own device affine already resolves bare dependencies. Corrects the earlier misdiagnosis recorded in the test comment; the write-then-replicated-read sequence is restored in the test. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Some PR to combine locality domain (#10658 ) and python support (#5315)
closes
Checklist