[STF] Add python bindings

AGENTS.md

@@ -207,6 +207,7 @@ Supported versions: `3.10`, `3.11`, `3.12`, `3.13`
 
 * **cuda.compute** — Device-level algorithms, iterators, custom GPU types
 * **cuda.coop._experimental** — Block/warp-level primitives for Numba CUDA
+* **cuda.stf._experimental** — Stream Task Flow (CUDASTF) Python bindings, shipped in the standalone `cuda-stf` package (Linux only)
 * **cuda.cccl.headers** — Programmatic access to headers
 
 ### Installation
@@ -262,14 +263,16 @@ include_paths = headers.get_include_paths()
 ./ci/test_cuda_coop_python.sh -py-version 3.10
 ./ci/test_cuda_cccl_headers_python.sh -py-version 3.10
 ./ci/test_cuda_cccl_examples_python.sh -py-version 3.10
+./ci/test_cuda_stf_python.sh -py-version 3.10  # Linux only
 ```
 
 Test organization:
 
 * `tests/compute` — Algorithms and iterators
 * `tests/coop` — Cooperative primitives
 * `tests/headers` — Header integration
-* `test_examples.py` — Runs compute/coop examples
+* `python/cuda_stf/tests/stf` — Sequential Task Flow (separate `cuda-stf` package, Linux only)
+* `test_examples.py` — Runs compute/coop examples (STF examples live in `python/cuda_stf/tests/test_examples.py`)
 
 ---
 

ci/build_cuda_stf_python.sh

@@ -0,0 +1,169 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ci_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+usage="Usage: $0 -py-version <python_version> [additional options...]"
+
+# shellcheck source=ci/util/python/common_arg_parser.sh
+source "$ci_dir/util/python/common_arg_parser.sh"
+parse_python_args "$@"
+
+# Check if py_version was provided (this script requires it)
+require_py_version "$usage" || exit 1
+
+echo "Docker socket: " "$(ls /var/run/docker.sock)"
+
+if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
+  # Prepare mount points etc for getting artifacts in/out of the container.
+  # shellcheck source=ci/util/artifacts/common.sh
+  source "$ci_dir/util/artifacts/common.sh"
+  # Note that these mounts use the runner (not the devcontainer) filesystem for
+  # source directories because of docker-out-of-docker quirks.
+  # The workflow-job GH actions make sure that they exist before running any
+  # scripts.
+  action_mounts=(
+    --mount "type=bind,source=${ARTIFACT_ARCHIVES},target=${ARTIFACT_ARCHIVES}"
+    --mount "type=bind,source=${ARTIFACT_UPLOAD_STAGE},target=${ARTIFACT_UPLOAD_STAGE}"
+  )
+else
+  # If not running in GitHub Actions, we don't need to set up artifact mounts.
+  action_mounts=()
+fi
+
+# cuda_stf must be built in a container that can produce manylinux wheels, and
+# has the CUDA toolkit installed. We use the rapidsai/ci-wheel image for this.
+# We build separate wheels using separate containers for each CUDA version,
+# then merge them into a single wheel. CUDASTF is Linux-only.
+
+readonly cuda12_version=12.9.1
+readonly cuda13_version=13.1.1
+readonly devcontainer_version=26.04
+readonly devcontainer_distro=rockylinux8
+# Use a baseline Python tag for the rapidsai ci-wheel image. The requested
+# py_version is installed inside the container by setup_python_env (uv).
+readonly devcontainer_python_version=3.10
+
+if [[ "$(uname -m)" == "aarch64" ]]; then
+  cuda12_image="rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda12_version}-${devcontainer_distro}-py${devcontainer_python_version}-arm64"
+  cuda13_image="rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda13_version}-${devcontainer_distro}-py${devcontainer_python_version}-arm64"
+else
+  cuda12_image="rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda12_version}-${devcontainer_distro}-py${devcontainer_python_version}"
+  cuda13_image="rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda13_version}-${devcontainer_distro}-py${devcontainer_python_version}"
+fi
+# shellcheck disable=SC2034
+readonly cuda12_image
+# shellcheck disable=SC2034
+readonly cuda13_image
+
+mkdir -p wheelhouse
+
+# Shared caches across the cu12 + cu13 wheel builds. Both jobs compile an
+# identical LLVM/clang tree (LLVM has no CUDA dep), so a shared ccache cuts
+# the second build's LLVM phase substantially; a shared CPM source cache skips
+# the second LLVM git clone entirely.
+mkdir -p ./.ccache ./.cpm-cache
+host_ccache_dir="${HOST_WORKSPACE:?}/.ccache"
+host_cpm_cache_dir="${HOST_WORKSPACE:?}/.cpm-cache"
+
+for ctk in 12 13; do
+  image="cuda${ctk}_image"
+  image="${!image}"
+  echo "::group::⚒️ Building CUDA $ctk cuda-stf wheel on $image"
+  (
+    set -x
+    docker pull "$image"
+    docker run --rm -i \
+        --workdir /workspace/python/cuda_stf \
+        --mount "type=bind,source=${HOST_WORKSPACE:?},target=/workspace/" \
+        --mount "type=bind,source=${host_ccache_dir},target=/root/.ccache" \
+        --mount "type=bind,source=${host_cpm_cache_dir},target=/root/.cpm-cache" \
+        "${action_mounts[@]}" \
+        --env "py_version=${py_version}" \
+        --env "GITHUB_ACTIONS=${GITHUB_ACTIONS:-}" \
+        --env "GITHUB_RUN_ID=${GITHUB_RUN_ID:-}" \
+        --env "JOB_ID=${JOB_ID:-}" \
+        --env "CCACHE_DIR=/root/.ccache" \
+        --env "CPM_SOURCE_CACHE=/root/.cpm-cache" \
+        "$image" \
+        /workspace/ci/build_cuda_stf_wheel.sh
+    # Prevent GHA runners from exhausting available storage with leftover images:
+    if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
+      docker rmi -f "$image"
+    fi
+  )
+  echo "::endgroup::"
+done
+
+echo "Merging CUDA wheels..."
+
+# Set up a Python environment for the merge/repair steps.
+source "$ci_dir/pyenv_helper.sh"
+setup_python_env "${py_version}"
+
+# Needed for unpacking and repacking wheels.
+python -m pip install wheel
+
+# Find the built wheels
+cu12_wheel=$(find wheelhouse -name "cuda_stf-*cu12*.whl" | head -1)
+cu13_wheel=$(find wheelhouse -name "cuda_stf-*cu13*.whl" | head -1)
+
+if [[ -z "$cu12_wheel" ]]; then
+  echo "Error: CUDA 12 cuda-stf wheel not found in wheelhouse/"
+  ls -la wheelhouse/
+  exit 1
+fi
+
+if [[ -z "$cu13_wheel" ]]; then
+  echo "Error: CUDA 13 cuda-stf wheel not found in wheelhouse/"
+  ls -la wheelhouse/
+  exit 1
+fi
+
+echo "Found CUDA 12 wheel: $cu12_wheel"
+echo "Found CUDA 13 wheel: $cu13_wheel"
+
+# Merge the wheels
+python python/cuda_stf/merge_cuda_wheels.py "$cu12_wheel" "$cu13_wheel" --output-dir wheelhouse_merged
+
+# Install auditwheel and repair the merged wheel
+python -m pip install patchelf auditwheel
+for wheel in wheelhouse_merged/cuda_stf-*.whl; do
+    echo "Repairing merged wheel: $wheel"
+    python -m auditwheel repair \
+        --exclude 'libnvrtc.so.12' \
+        --exclude 'libnvrtc.so.13' \
+        --exclude 'libnvJitLink.so.12' \
+        --exclude 'libnvJitLink.so.13' \
+        --exclude 'libcudart.so.12' \
+        --exclude 'libcudart.so.13' \
+        --exclude 'libcuda.so.1' \
+        "$wheel" \
+        --wheel-dir wheelhouse_final
+done
+
+# Clean up intermediate files and move only the final merged wheel to wheelhouse
+rm -rf wheelhouse/*  # Clean existing wheelhouse
+mkdir -p wheelhouse
+
+# Move only the final repaired merged wheel
+if ls wheelhouse_final/cuda_stf-*.whl 1> /dev/null 2>&1; then
+    mv wheelhouse_final/cuda_stf-*.whl wheelhouse/
+    echo "Final merged wheel moved to wheelhouse"
+else
+    echo "No final repaired wheel found, moving unrepaired merged wheel"
+    mv wheelhouse_merged/cuda_stf-*.whl wheelhouse/
+fi
+
+# Clean up temporary directories
+rm -rf wheelhouse_merged wheelhouse_final
+
+echo "Final wheels in wheelhouse:"
+ls -la wheelhouse/
+
+if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
+  # Upload under a distinct artifact name so it does not clobber the cuda-cccl
+  # wheel (both build jobs run in project 'python').
+  wheel_artifact_name="$(CCCL_WHEEL_KIND=stf ci/util/workflow/get_wheel_artifact_name.sh)"
+  ci/util/artifacts/upload.sh "$wheel_artifact_name" 'wheelhouse/.*'
+fi

ci/build_cuda_stf_wheel.sh

@@ -0,0 +1,82 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Target script for `docker run` command in build_cuda_stf_python.sh
+# The /workspace pathnames are hard-wired here.
+
+# Install GCC 13 toolset (needed for the build) and ccache (shared between
+# cu12 and cu13 builds via /root/.ccache bind-mount from the host).
+/workspace/ci/util/retry.sh 5 30 dnf -y install \
+  gcc-toolset-13-gcc gcc-toolset-13-gcc-c++ ccache
+
+# When the caller bind-mounts a ccache dir, wire it through to CMake. This
+# transparently caches every compile, so the second wheel build (cu13 after
+# cu12, or vice versa) reuses the entire LLVM/clang object tree.
+if [[ -n "${CCACHE_DIR:-}" ]]; then
+  export CMAKE_C_COMPILER_LAUNCHER=ccache
+  export CMAKE_CXX_COMPILER_LAUNCHER=ccache
+  export CMAKE_CUDA_COMPILER_LAUNCHER=ccache
+  echo "ccache enabled: CCACHE_DIR=${CCACHE_DIR}"
+  ccache --version 2>&1 | head -1 || true
+  ccache --show-stats 2>&1 | head -5 || true
+fi
+echo -e "#!/usr/bin/env bash\nsource /opt/rh/gcc-toolset-13/enable" >/etc/profile.d/enable_devtools.sh
+# shellcheck disable=SC1091
+source /etc/profile.d/enable_devtools.sh
+
+# Check what's available
+command -v gcc
+gcc --version
+command -v nvcc
+nvcc --version
+
+# Set up Python environment
+# shellcheck source=ci/pyenv_helper.sh
+source /workspace/ci/pyenv_helper.sh
+# shellcheck disable=SC2154
+setup_python_env "${py_version}"
+command -v python
+python --version
+echo "Done setting up python env"
+
+# Figure out the version to use for the package, we need repo history
+if "$(git rev-parse --is-shallow-repository)"; then
+  git fetch --unshallow
+fi
+# Match the cuda-cccl version prefix so cuda-stf and cuda-cccl stay in lockstep.
+export PACKAGE_VERSION_PREFIX="0.1."
+package_version=$(/workspace/ci/generate_version.sh)
+echo "Using package version ${package_version}"
+# Override the version used by setuptools_scm to the custom version
+export SETUPTOOLS_SCM_PRETEND_VERSION_FOR_CUDA_STF="${package_version}"
+
+cd /workspace/python/cuda_stf
+
+# Determine CUDA version from nvcc
+cuda_version=$(nvcc --version | grep -oP 'release \K[0-9]+\.[0-9]+' | cut -d. -f1)
+echo "Detected CUDA version: ${cuda_version}"
+
+# Configure compilers:
+CXX="$(command -v g++)"
+export CXX
+CUDACXX="$(command -v nvcc)"
+export CUDACXX
+CUDAHOSTCXX="$(command -v g++)"
+export CUDAHOSTCXX
+
+# Build the wheel
+python -m pip wheel --no-deps --verbose --wheel-dir dist .
+
+# Rename wheel to include CUDA version suffix
+for wheel in dist/cuda_stf-*.whl; do
+    if [[ -f "$wheel" ]]; then
+        base_name=$(basename "$wheel" .whl)
+        new_name="${base_name}.cu${cuda_version}.whl"
+        mv "$wheel" "dist/${new_name}"
+        echo "Renamed wheel to: ${new_name}"
+    fi
+done
+
+# Move wheel to output directory
+mkdir -p /workspace/wheelhouse
+mv dist/cuda_stf-*.cu*.whl /workspace/wheelhouse/

ci/matrix.yaml

@@ -91,6 +91,13 @@ workflows:
     - {jobs: ['test'], project: 'python', ctk: ['12.X',        '13.X'], py_version: ['3.10'], gpu: 'l4', cxx: ['gcc13', 'msvc2022']}
     - {jobs: ['test'], project: 'python', ctk: ['12.X','13.0', '13.X'], py_version: ['3.14'], gpu: 'l4', cxx: ['gcc13', 'msvc2022']}
     - {jobs: ['test'], project: 'python', py_version: '3.14', gpu: 'h100', cxx: 'gcc13'}
+    # py3.13 base wheel is needed by the cuda.stf._experimental tests below,
+    # which depend on `build_py_wheel`. Linux/gcc13 only, matching the STF
+    # matrix footprint.
+    - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.X'], py_version: ['3.13'], gpu: 'l4', cxx: 'gcc13'}
+    # cuda.stf._experimental -- pinned to gcc13, Linux only
+    - {jobs: ['test_py_stf'], project: 'python', ctk: ['12.X', '13.X'], py_version: ['3.13'], gpu: 'l4', cxx: 'gcc13'}
+    - {jobs: ['test_py_stf'], project: 'python', ctk: '13.X', py_version: '3.13', gpu: 'h100_2gpu', cxx: 'gcc13'}
     - {jobs: ['test_py_compute_minimal'], project: 'python', ctk: '13.X', py_version: '3.14', gpu: 'l4', cxx: 'gcc13'}
     # CCCL packaging:
     - {jobs: ['test'], project: 'packaging', ctk: '12.0', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080', args: '-min-cmake'}
@@ -152,6 +159,7 @@ workflows:
     - {project: 'cccl_c_parallel', jobs: ['test'], ctk: '13.X', cxx: 'gcc13', gpu: 'rtxpro6000', sm: 'gpu'}
     - {project: 'cccl_c_stf',      jobs: ['test'], ctk: '13.X', cxx: 'gcc13',               gpu: 'rtx2080', sm: 'gpu'}
     - {project: 'python', jobs: ['test'], ctk: '13.X', py_version: '3.14', gpu: 'l4', cxx: ['gcc13', 'msvc2022']}
+    - {project: 'python', jobs: ['test_py_stf'], ctk: '13.X', py_version: '3.13', gpu: 'l4', cxx: 'gcc13'}
     # Packaging / install
     - {project: 'packaging', jobs: ['test'], ctk: '13.X', cxx: ['gcc', 'clang'], gpu: 'rtx2080', sm: 'gpu'}
     - {project: 'packaging', jobs: ['test'], args: '-min-cmake', gpu: 'rtx2080', sm: 'gpu'}
@@ -230,6 +238,9 @@ workflows:
     # Python -- pinned to gcc13 / msvc2022 on Linux for consistency across CTK images
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.0', '13.X'], py_version: ['3.10', '3.11', '3.12', '3.13', '3.14'], gpu: 'l4', cxx: ['gcc13', 'msvc2022']}
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.X'], py_version: '3.14', gpu: 'h100', cxx: 'gcc13'}
+     # cuda.stf._experimental -- pinned to gcc13, Linux only
+    - {jobs: ['test_py_stf'], project: 'python', ctk: ['12.X', '13.X'], py_version: ['3.10', '3.13'], gpu: 'l4', cxx: 'gcc13'}
+    - {jobs: ['test_py_stf'], project: 'python', ctk: '13.X', py_version: '3.13', gpu: 'h100', cxx: 'gcc13'}
     # CCCL packaging:
     - {jobs: ['test'], project: 'packaging', ctk: '12.0', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080', args: '-min-cmake'}
     - {jobs: ['test'], project: 'packaging', ctk: '12.X', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080'}
@@ -325,6 +336,9 @@ workflows:
     # Python -- pinned to gcc13 / msvc2022 for consistency across CTK images
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.0', '13.X'], py_version: ['3.10', '3.11', '3.12', '3.13', '3.14'], gpu: 'l4', cxx: ['gcc13', 'msvc2022']}
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.X'], py_version: '3.14', gpu: 'h100', cxx: ['gcc13', 'msvc2022']}
+    # cuda.stf._experimental -- pinned to gcc13, Linux only
+    - {jobs: ['test_py_stf'], project: 'python', ctk: ['12.X', '13.X'], py_version: ['3.10', '3.13'], gpu: 'l4', cxx: 'gcc13'}
+    - {jobs: ['test_py_stf'], project: 'python', ctk: '13.X', py_version: '3.13', gpu: 'h100', cxx: 'gcc13'}
     # CCCL packaging:
     - {jobs: ['test'], project: 'packaging', ctk: '12.0', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080', args: '-min-cmake'}
     - {jobs: ['test'], project: 'packaging', ctk: '12.X', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080'}
@@ -356,7 +370,7 @@ workflows:
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.0', '13.X'], py_version: ['3.10', '3.11', '3.12', '3.13', '3.14'], gpu: 'l4', cxx: ['gcc13', 'msvc2022']}
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.X'], py_version: '3.14', gpu: 'h100', cxx: ['gcc13', 'msvc2022']}
     - {jobs: ['test'], project: 'python', cpu: 'arm64', ctk: ['12.X', '13.X'], py_version: ['3.10', '3.11', '3.12', '3.13', '3.14'], gpu: 'l4', cxx: 'gcc13'}
-
+    - {jobs: ['test_py_stf'], project: 'python', ctk: ['12.X', '13.X'], py_version: ['3.10', '3.13'], gpu: 'l4', cxx: 'gcc13'}
 
   # This is just used to ensure that we generate devcontainers for all images we build.
   # These do not map to any actual jobs.
@@ -384,6 +398,8 @@ workflows:
   exclude:
     # GPU runners are not available on Windows.
     - {jobs: ['test', 'test_gpu', 'test_nolid', 'test_lid0', 'test_lid1', 'test_lid2'], cxx: ['msvc2019', 'msvc14.39', 'msvc2022', 'msvc2026']}
+    # cuda.stf._experimental is Linux-only.
+    - {jobs: ['test_py_stf'], cxx: ['msvc2019', 'msvc14.39', 'msvc2022', 'msvc2026']}
     # cudax doesn't support C++17 on msvc:
     - {project: 'cudax', std: 17, cxx: ['msvc2019', 'msvc14.39', 'msvc2022', 'msvc2026']}
 
@@ -544,11 +560,17 @@ jobs:
 
   # Python:
   build_py_wheel:   { name: "Build cuda.cccl",             gpu: false, invoke: { prefix: 'build_cuda_cccl'} }
+  build_py_stf_wheel: { name: "Build cuda.stf",            gpu: false, invoke: { prefix: 'build_cuda_stf'} }
   test_py_headers:  { name: "Test cuda.cccl.headers",      gpu: true,  needs: 'build_py_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_cccl_headers'} }
   test_py_coop:     { name: "Test cuda.coop._experimental",  gpu: true,  needs: 'build_py_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_coop'} }
   test_py_par:      { name: "Test cuda.compute",     gpu: true,  needs: 'build_py_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_compute'} }
   test_py_compute_minimal: { name: "Test cuda.compute minimal", gpu: true, needs: 'build_py_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_compute_minimal'} }
   test_py_examples: { name: "Test cuda.cccl.examples",     gpu: true,  needs: 'build_py_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_cccl_examples'} }
+  # cuda-stf is standalone: it ships the STF bindings + its own headers and has
+  # no hard cuda-cccl dependency, so only the cuda-stf wheel is needed. The
+  # interop tests' optional cuda-cccl (for cuda.compute) comes from the test
+  # extra.
+  test_py_stf:      { name: "Test cuda.stf._experimental", gpu: true,  needs: 'build_py_stf_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_stf'} }
 
   # Run jobs for 'target' project (ci/util/build_and_test_targets.sh):
   run_cpu: { gpu: false }

ci/project_files_and_dependencies.yaml

@@ -148,11 +148,15 @@ projects:
   python:
     name: "Python"
     matrix_project: "python"
-    lite_dependencies: [cccl_c_parallel_public]
+    lite_dependencies: [cccl_c_parallel_public, cccl_c_stf]
     full_dependencies: []
     include_regexes:
-    - "python/"
+    - "python/cuda_cccl/"
+    # cuda-stf is a separate distribution but is built/tested by the same
+    # 'python' matrix jobs (build_py_stf_wheel / test_py_stf).
+    - "python/cuda_stf/"
     - "pyproject.toml"
+    exclude_regexes: []
 
   packaging:
     name: "CCCL Packaging"

ci/test/inspect_changes/c2h_dependency.output

@@ -1,2 +1,2 @@
 FULL_BUILD=tidy
-LITE_BUILD=libcudacxx cub cudax cccl_c_parallel cccl_c_parallel_v2 python_v2 cccl_c_stf packaging
+LITE_BUILD=libcudacxx cub cudax cccl_c_parallel cccl_c_parallel_v2 python_v2 cccl_c_stf python packaging