[STF] Add structured tensor partitions and placement evaluation to the places layer#9804
Merged
Conversation
Contributor
Closed
2 tasks
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds CuTe-style tensor partitioning, geometry-aware localized allocation, STF dimension mapping, a partitioned AXPY example, and placement and residency tests. ChangesPartitioned placement and allocation
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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: 714eb680-5a6a-4dd1-a6c6-a6499855b929
📒 Files selected for processing (1)
cudax/examples/stf/partitioned_axpy.cu
Contributor
Author
|
/ok to test 4aba5e3 |
This comment has been minimized.
This comment has been minimized.
Contributor
Author
|
CI caught the one in-tree consumer of the removed byte-count-allocate-on-composite path: the |
Contributor
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: a14377dc-58fb-49a0-b8b3-dd4539132e09
📒 Files selected for processing (2)
cudax/examples/places/thrust_device_data_place_allocator.cudocs/cudax/places.rst
Contributor
Author
|
/ok to test 6f84405 |
This comment has been minimized.
This comment has been minimized.
2 tasks
Contributor
Author
|
Follow-up stacked PR: #9808 completes the compute side — the same |
Jacobfaib
approved these changes
Jul 13, 2026
caugonnet
force-pushed
the
stf-cute-partitions
branch
from
July 16, 2026 07:26
b3d0fab to
9896ffa
Compare
2 tasks
Composite data places distribute an allocation according to a
partitioner that maps element coordinates to places, but the raw
allocation entry only carried a byte count: the partitioner received
flat byte offsets with fabricated 1-D extents, so any coordinate-based
mapper silently placed everything on one device. The tensor geometry
must reach the allocation seam.
Add a second allocate(data_dims, elemsize, stream) virtual to
data_place_interface with a base default forwarding to the byte-count
form: regular places inherit it unchanged ("same code, different
places" now extends to geometry-aware callers), while composite places
override it with the real localized_array path and reject the
byte-count form, whose implicit 1-D mode is removed. dim4 gains
index_to_pos(), the inverse of get_index(), shared by the runtime
slice interfaces and the raw composite path so the two linearization
conventions cannot drift.
Also make the block-owner decision reproducible (one seeded generator
per placement instead of a fresh random_device per block, probe count
now a parameter) and return placement statistics as a struct
(localized_stats: bytes per place, sampling accuracy - previously
stderr-only under CUDASTF_LOCALIZED_ALLOC_STATS, whose output is
preserved). A new evaluate_localized_placement() entry runs the exact
same decision procedure without allocating, so a candidate mapping can
be scored and tuned before committing memory.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Add a structured way to express how a tensor is distributed over a
grid of places: a per-dimension specification in the style of JAX's
PartitionSpec ("dimension 1, blocked over grid axis 0") lowered by
make_partition() onto a CuTe-style two-mode layout of flattened
(extent, stride) leaves.
Split dimensions are padded up to divisibility - the CuTe predication
idiom (partition the rounded shape, predicate against the true
extents) - which makes the layout exact and bijective over the padded
space. Exactness is what keeps the machinery trivial: validation is an
O(leaves) sorted-stride mixed-radix check, ownership queries are a
closed-form divmod chain, and no CUTLASS/CuTe dependency is needed
since only that subset of the layout algebra is exercised. Accessors
(dims, leaves, per-place offsets) expose everything a front-end needs
to derive further artifacts from the same object that drives data
placement.
The partition feeds the existing localized allocation machinery as a
stateful owner function: evaluate_localized_placement() accepts it
directly, and make_composite_data_place() wraps it as a composite data
place (per-tensor by nature, since a padded partition is specific to
one tensor's extents; the block-majority engine still decides where
each block lives). Routing such places through the runtime logical
data path is a recorded follow-up.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
placement.cu covers the evaluator (block-aligned split, majority tie-breaking on a straddling block, seeded determinism, cute partition agreeing with the equivalent classic partitioner), geometry-aware allocation on both composite flavors including kernel write/readback, and the rejection paths (byte-count allocate on a composite, extent mismatch with a partition). With two or more devices it additionally asserts the physical residency of each block on its owner and exercises the peer-access mappings; those checks self-skip on single-GPU machines so the multi-GPU CI provides the real coverage. cute_partition.cuh carries co-located unittests for the builders (including the uneven-extent padding case that motivates the padded design), owner sweeps, and validation rejections. The partitioned_axpy example shows the intended workflow: express the partition once, score it with evaluate_localized_placement() before allocating, run STF tasks over logical data on a composite place, and perform a raw geometry-aware allocation outside any STF context. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The bytes-per-place map is keyed by place strings, which cannot cross a C FFI boundary usefully. Also record bytes per linear grid index so bindings can return the distribution as a plain array. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Overloading allocate() made every data_place_interface subclass that overrides the byte-count form partially override the virtual set, tripping -Woverloaded-virtual until it adds a using-declaration - a recurring tax on every current and future place implementation. A distinct name removes the whole problem class (no name hiding, no overload-ambiguity concerns) and aligns with the C API, which cannot overload and needs a distinct name anyway. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Adversarial review fixes, none changing the intended semantics: - Guard rank-0 (scalar) slices in the composite delinearize: their data_dims report an extent of 0, so index_to_pos would divide by zero where the previous per-rank lambda returned the origin. - Validate elemsize (>= 1, <= block size) at the placement seams instead of dividing by zero or silently sampling garbage when a block holds no whole element. - Compute each block's first element exactly (i * block_size / elemsize, as before the refactor) so non-dividing element sizes do not accumulate a sampling-window drift across blocks. - Check the partition's grid extents against the execution place grid (cute composite construction and evaluation): a mismatched pair previously mis-attributed bytes or aborted deep in get_place. - Reject zero-extent leaves and grid axes, and out-of-range place_offset indices, instead of wrapping silently. - Compare data_dims for real in the allocation-cache equality (was assert-only: transposed extents could produce a false cache hit in release builds). - Make the straddling-block test structural (62.5/37.5 split: the all-probes-in-majority probability is ~1e-13 for any uniform sampler) instead of relying on the stdlib's implementation-defined distribution sequence. - Fall back to a single device in the partitioned AXPY example when peer access is unavailable (the single task touches every page from each device), and factor the default block-size selection into one helper. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
localized_array already validates VMM support with a clear error, and no other composite-place example performs its own check; the raw driver boilerplate distracted from the example's point. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ample exec_place::all_devices() already enumerates the devices, and the places machine sets up peer and mempool access between them - the example re-implemented both, plus a peer-access fallback that only existed because it launched a single kernel over the whole composite buffer from one device. Follow the idiomatic grid-task pattern of 01-axpy-places instead (each place computes its own blocked chunk on its own stream), which needs no cross-device access at all. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
With ceil-division chunks, i * chunk can exceed N when there are more places than chunks; the unsigned N - start then underflows and the kernel writes out of bounds. Unreachable with the example's sizes, but the loop demonstrates the canonical pattern and must be safe in general. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The example's memory resource byte-allocated on a composite place, which now throws (a byte count alone cannot carry the geometry the partitioner needs). A memory resource genuinely hands out untyped bytes, so it states that explicitly with allocate_nd(dim4(bytes), 1) - equivalent for every other place type, and byte-granular distribution for composites, matching the previous behavior. Document allocate_nd and the composite requirement in the places guide. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replace the std::vector leaf storage of cute_partition with cuda::std::array (bounded by a documented max_leaves = 16 per mode; make_partition emits at most 2 per dimension). The partition becomes trivially copyable and its queries (owner, dims, leaf spans) are host/device callable, so it can cross the kernel boundary by value - the property a future parallel_for integration needs, with no separate device-side snapshot type. Accessors now return cuda::std::span over the filled prefix; construction keeps its std::vector convenience signature and rejects leaf counts beyond the capacity. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
make_partition and its surroundings introduce several non-obvious concepts that only lived in code comments: the reference-shape binding (and how it differs from the scale-free classic policies), padding to divisibility and the predication idiom, evaluation before allocation, the per-tensor nature of partition-backed composite places, and the linearization convention. Give them a section in the places guide, anchored against the existing partitioning-policies documentation. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
allocate_nd(dim4(nbytes), 1) only works for composites built from scale-free partitioners; a structured-partition composite validates the extents against the partition's tensor and rejects a flat byte count. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Cedric AUGONNET <caugonnet@nvidia.com>
sin/cos/fabs were relying on a transitive include. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Cedric AUGONNET <caugonnet@nvidia.com>
…tion Reversing only the extents would re-target each dim_spec at the wrong axis; the spec list and owner() coordinates reverse together. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Cedric AUGONNET <caugonnet@nvidia.com>
Describe the API directly as a structured per-dimension partition specification to avoid implying unsupported JAX semantics.
Without replication, an unbound grid axis pins every element to coordinate 0 on that axis and leaves the other places idle. Fail at construction time so callers must collapse the grid or bind every axis to a tensor dimension.
caugonnet
force-pushed
the
stf-cute-partitions
branch
from
July 16, 2026 14:26
a7509f1 to
da9b920
Compare
Contributor
Author
|
/ok to test a7509f1 |
Contributor
@caugonnet, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
Contributor
Author
|
/ok to test d98448a |
Contributor
🥳 CI Workflow Results🟩 Finished in 1h 08m: Pass: 100%/57 | Total: 1d 07h | Max: 1h 08m | Hits: 12%/200595See results here. |
caugonnet
added a commit
to caugonnet/cccl
that referenced
this pull request
Jul 17, 2026
The C++ runtime partition class merged through NVIDIA#9804/NVIDIA#9808 as cute_partition_descriptor (cute_partition now names the typed template), and the runtime factory as make_partition_descriptor. Constructor and factory signatures are unchanged. 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
This adds the data-side machinery for distributing a tensor over a grid of places, designed around the existing
localized_arrayVMM allocator and its block-majority placement engine:cute_partition(__places/cute_partition.cuh): a structured description of a tensor partition — a JAX-like per-dimension specification (make_partition(dims, {dim_spec{dim_policy::blocked, /*axis*/ 0}}, grid_dims), i.e. "dimension 1, blocked over grid axis 0") lowered onto a CuTe-style two-mode strided layout. Split dimensions are padded up to divisibility (the CuTe predication idiom: partition the rounded shape, coordinates beyond the true extents own no bytes), which makes the layout exact and bijective over the padded space: validation is an O(leaves) sorted-stride mixed-radix check, ownership queries are a closed-form divmod chain, and no CUTLASS dependency is needed. Accessors (dims, leaves, per-place offsets) expose everything a front-end needs to derive further artifacts from the same object that drives data placement.data_place::allocate_nd(data_dims, elemsize, stream): a second virtual with a base default forwarding to the byte-count form. Regular places inherit it unchanged; composite places override it so the partitioner finally receives true element coordinates. The byte-countallocate()on composite places now throws: it previously fed the mapper fabricated 1-D extents, silently placing everything on one device for any coordinate-based partitioner (behavior change: the raw byte path on composites was unusable-in-disguise; the runtime logical-data path is unaffected).dim4gainsindex_to_pos(), shared by the runtime slice interfaces and the raw composite path so the linearization conventions cannot drift.evaluate_localized_placement()+localized_stats: run the exact same (now seeded, reproducible) block-owner decision procedure without allocating, returning the statistics previously only printed to stderr underCUDASTF_LOCALIZED_ALLOC_STATS(bytes per place/grid position, sampling accuracy) — so a candidate mapping can be scored and tuned before committing memory.Tests: co-located unittests for the partition algebra (padding of uneven extents, owner sweeps against
blocked_partition_custom, validation rejections), aplacement.cuGPU test covering the evaluator, both composite flavors, and the rejection paths, plus a multi-GPU tier (physical block residency on its owner, peer-access path) that self-skips below two devices. Apartitioned_axpyexample shows the intended workflow: express the partition once, score it, run STF tasks over composite-placed logical data, and perform a raw geometry-aware allocation.Follow-ups deliberately out of scope: routing
cute_partition-backed places through the runtime logical-data path, a layout-driven partitioner forparallel_for, RAII rollback for mid-constructor VMM failures (pre-existing).Checklist
🤖 Generated with Claude Code