[STF] pytorch.localized with locality domains - #10705
Draft
caugonnet wants to merge 807 commits into
Draft
Conversation
Validate Burger reference step grouping so environment overrides cannot silently drop tail steps, and clean up minor stale test code noted in review.
Tighten STF Python lifecycle cleanup, lazy imports, binding shims, CI diagnostics, and API docs based on the comprehensive follow-up review.
Describe the Burger reference example as a PyTorch baseline for the same solver rather than explaining prior planning context.
green_ctx_data_place_impl had no mem_create override, so it inherited the data_place_interface default returning CUDA_ERROR_NOT_SUPPORTED. A composite/ localized_array allocation whose affine places are green contexts therefore failed at cuMemCreate. Add a device-style mem_create (PINNED, DEVICE, location.id = view_.devid) mirroring data_place_device, so each partition gets its own physical VMM chunk. Also include the CUgreenCtx handle in to_string() so distinct contexts no longer render identically as "green_ctx(dev=0)" -- the localized allocation stats now report them as separate places. Assisted-by: Cursor (Claude Opus 4.8) Signed-off-by: Cedric AUGONNET <caugonnet@nvidia.com>
Reuse the existing cuda.compute stream protocol handling in STF Python bindings so stream arguments honor __cuda_stream__ objects while preserving raw pointer support.
Keep CUDA Array Interface descr metadata when STF exports structured task arguments, and centralize STF CAI dtype parsing to avoid losing field layout information.
Destroy the opaque exec_place_resources handle returned by ctx.place_resources so the wrapper is released while the context keeps owning the underlying stream pools.
Raise Python ValueError when a task or host launch receives logical data from another context, avoiding later C++ aborts on context mismatches.
Run cuda.stf._experimental Python tests on the h100_2gpu runner so PR CI covers multi-GPU STF behavior.
Move the Linux-only cuda.stf._experimental bindings out of the cuda-cccl
wheel into a new, optional python/cuda_stf distribution, following the
split-package model used by cuda-coop. The runtime import path
cuda.stf._experimental is preserved; users now install cuda-stf[cu12|cu13].
- Move cuda/stf source and STF tests to python/cuda_stf, dropping the empty
cuda/stf/__init__.py so cuda/cuda.stf/cuda.cccl remain shared namespaces.
- Add cuda-stf packaging (pyproject, CMakeLists, README, LICENSE, wheel
merger). cuda-stf depends on cuda-cccl for headers and CUDA-version
detection, and ships only the cudax/C-STF headers cuda-cccl does not
provide to avoid file-ownership conflicts.
- Remove STF build/install and wheel-merge behavior from cuda-cccl and drop
STF from its examples runner.
- Wire CI: build_cuda_stf_{python,wheel}.sh, a build_py_stf_wheel job, a
distinct STF wheel artifact (CCCL_WHEEL_KIND), and a test_py_stf lane that
installs both wheels and runs python/cuda_stf/tests.
- Update docs and AGENTS.md to document the cuda-stf install path.
cuda-stf previously required cuda-cccl at runtime for CUDA version detection (cuda.cccl._cuda_version_utils) and header discovery (cuda.cccl.headers). This copies the tiny version-detection helper into cuda-stf, ships STF's own cudax + C-STF headers under cuda/stf/_experimental/include, and makes cuda-cccl an optional peer used only to supply libcudacxx/CUB/Thrust include paths for external C++ compilation. cuda-cccl remains in the test extras for the cuda.compute interop tests. This also fixes the broken CI matrix parse: test_py_stf now depends on a single producer (build_py_stf_wheel) instead of a list of producers, which the workflow-build machinery does not support (TypeError: unhashable type: 'list').
(cherry picked from commit 9d37b5e)
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>
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>
…d 'origin/stf-cute-partitions-python' into x-9892-integration
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>
… x-9892-integration
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>
… x-9892-integration
…ntyped access) Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… x-9892-integration
…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>
… x-9892-integration
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>
… x-9892-integration
The controlled experiments behind the 2026-08-08 GB300 findings, as one reproducible script (benchmarks/localization_ab.py): streaming fused compiled pointwise (vanilla loses under caps: near/vanilla 1.27x at 250-400 W, wins 1.22x at the wall), DRAM- and L2-resident random gathers (placement-neutral at nn concurrency), and sustained-power energy (near ~0.95x J/GB uncapped). Requires the locality-domain bindings and skips cleanly without them; confinement via exec_place.pick_stream domain streams; power caps external (powerbot). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… x-9892-integration
caugonnet
commented
Aug 7, 2026
| """Stock torch.compile when triton is available, eager otherwise -- | ||
| map treats both identically (any callable).""" | ||
| try: | ||
| import triton # noqa: F401, PLC0415 |
Contributor
Author
There was a problem hiding this comment.
This probably makes more sense to skip the test entirely if triton is not available using pytest ?
caugonnet
commented
Aug 7, 2026
| return eg.machine(granularity="locality_domain") | ||
| if hasattr(eg, "locality_domains"): | ||
| return eg.locality_domains(0) | ||
| pytest.skip("locality-domain bindings not available (PR #10703)") |
Contributor
Author
There was a problem hiding this comment.
the PR name should ultimately disappear
Expose checked grid reshaping and contiguous-axis collapse through the C and Python APIs. Keep transformed grid handles independently owned and convert invalid shape errors into clean FFI failures instead of process termination.
A grid_dims product that does not match the number of places now raises a ValueError in the Python binding (stf_c_api review change) instead of a RuntimeError from the C layer. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Follow the C-order Python contract from the structured-partition work: reshape() takes a public C-order shape (linear place order is preserved because the public C-order enumeration and the native dimension-0-fastest enumeration coincide after reversal), collapse_axes() takes public axes mapped onto the native inclusive range [rank-1-last, rank-1-first] and validated against the grid's rank, and both stamp the result's grid rank so dims stays rank-aware. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…e bindings assume The reshape/collapse bindings cherry-picked from NVIDIA#9899 were written on top of the C-order public-shape contract, whose plumbing lived in the stf-cute-partitions-python branch. Hoist that slice here, matching its text exactly so the later NVIDIA#9892 merge is trivial: - _validate_extents/_fill_dim4_c_order/_native_to_public helpers (the one place the public C order and the native dimension-0-fastest ABI meet); - exec_place.dims returns a C-order tuple of the grid rank (scalar: (1,)); - exec_place_grid tracks _grid_rank (grid_rank property; set by create/ reshape/collapse_axes, 1-D creation paths default to 1); - create(grid_dims=) validates through the shared helpers; - task tracks the rank from set_exec_place; get_grid_dims returns the C-order shape; get_stream_ptrs sizes from the rank-length tuple; - test expectations updated to the contract ((2, 2) grids, scalar (1,)). The C++ reshape/collapse machinery itself is already in main (with _CCCL_HOST_API annotations), so the NVIDIA#9899 C++ commit was dropped as redundant during the fold. Co-Authored-By: Claude Fable 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
2 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
Experiments with locality domains with in pytorch.localized. This PR combines #9892 and #10703 (and by transitivity #5315 and #10658)
closes
Checklist