perf: give each warp lane its own groupby shared-memory aggregation slot - #3
Open
felipeblazing wants to merge 2 commits into
Open
Conversation
single_pass_shmem_aggs_kernel replicates its shared memory aggregation slots to
spread out the atomic updates, but the replication factor was a fixed
`32 / cardinality`. For a block of GROUPBY_BLOCK_SIZE=128 threads that is only 32
slots, so several threads of every warp still target the same slot.
That is much more expensive than the usual "atomics get replayed" intuition,
because a shared memory atomicAdd on a 64-bit type is not a native instruction.
The kernel's SASS contains no ATOMS.ADD at all -- every shared aggregation lowers
to ATOMS.CAST.SPIN.64, a lock/retry loop -- so colliding threads spin instead of
being serialized once by the hardware, and the cost grows far worse than linearly
in the number of colliding lanes.
Raise the factor until each lane of a warp owns a replica, after which no two
threads of a warp can collide. The room is already paid for: the launch reserves
enough shared memory for GROUPBY_CARDINALITY_THRESHOLD distinct groups, so a
low-cardinality block leaves most of that reservation idle while still paying its
occupancy cost. This spends it on replicas instead. Occupancy is unchanged.
Measured on GB300 / CUDA 13.2, cudf::groupby over 200M rows with two 1-char keys
and 9 output columns (8 FLOAT64 SUM + COUNT), whole-groupby wall time:
distinct groups | before | after |
4 | 86.2 ms | 19.6 ms | 4.39x
16 | 224.9 ms | 64.4 ms | 3.49x
64 | 346.6 ms | 204.3 ms | 1.70x
127 | 101.1 ms | 100.0 ms | 1.01x (factor is 1 either way)
The gain is a cliff, not a curve: sweeping the slot budget at 4 groups gives
84.9 ms at a factor of 16 and 19.8 ms at 32, then 19.7/19.7/19.6 ms at 64/128/256.
That is exactly the warp width, which is why the target is one replica per lane
and no more -- asking for more only costs initialization and merge work.
The factor is bounded so it cannot regress the previous rule: it never exceeds
what keeps every column resident (so the kernel still makes a single pass over the
input), never exceeds one replica per lane, never grows past a fraction of the
rows the block actually processes, and is never below the old fixed value. When
shared memory or cardinality leaves no room to grow it returns the old value
unchanged, which is the 127-group row above. It is computed once per block by
thread 0 rather than redundantly per thread.
Validated by comparing aggregate results against the previous factor at 1, 2, 3,
4, 7, 16, 33, 64 and 127 distinct groups, using integer-valued inputs so the sums
are exact in double and independent of summation order: identical on every group
and every column, with identical output row counts.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
mike-wendt
force-pushed
the
release/26.06
branch
5 times, most recently
from
August 21, 2026 19:39
eb2e22f to
d49c634
Compare
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
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.
single_pass_shmem_aggs_kernel replicates its shared memory aggregation slots to
spread out the atomic updates, but the replication factor was a fixed
32 / cardinality. For a block of GROUPBY_BLOCK_SIZE=128 threads that is only 32slots, so several threads of every warp still target the same slot.
That is much more expensive than the usual "atomics get replayed" intuition,
because a shared memory atomicAdd on a 64-bit type is not a native instruction.
The kernel's SASS contains no ATOMS.ADD at all -- every shared aggregation lowers
to ATOMS.CAST.SPIN.64, a lock/retry loop -- so colliding threads spin instead of
being serialized once by the hardware, and the cost grows far worse than linearly
in the number of colliding lanes.
Raise the factor until each lane of a warp owns a replica, after which no two
threads of a warp can collide. The room is already paid for: the launch reserves
enough shared memory for GROUPBY_CARDINALITY_THRESHOLD distinct groups, so a
low-cardinality block leaves most of that reservation idle while still paying its
occupancy cost. This spends it on replicas instead. Occupancy is unchanged.
Measured on GB300 / CUDA 13.2, cudf::groupby over 200M rows with two 1-char keys
and 9 output columns (8 FLOAT64 SUM + COUNT), whole-groupby wall time:
The gain is a cliff, not a curve: sweeping the slot budget at 4 groups gives
84.9 ms at a factor of 16 and 19.8 ms at 32, then 19.7/19.7/19.6 ms at 64/128/256.
That is exactly the warp width, which is why the target is one replica per lane
and no more -- asking for more only costs initialization and merge work.
The factor is bounded so it cannot regress the previous rule: it never exceeds
what keeps every column resident (so the kernel still makes a single pass over the
input), never exceeds one replica per lane, never grows past a fraction of the
rows the block actually processes, and is never below the old fixed value. When
shared memory or cardinality leaves no room to grow it returns the old value
unchanged, which is the 127-group row above. It is computed once per block by
thread 0 rather than redundantly per thread.
Validated by comparing aggregate results against the previous factor at 1, 2, 3,
4, 7, 16, 33, 64 and 127 distinct groups, using integer-valued inputs so the sums
are exact in double and independent of summation order: identical on every group
and every column, with identical output row counts.
Co-Authored-By: Claude Opus 5 noreply@anthropic.com