Skip to content

perf: give each warp lane its own groupby shared-memory aggregation slot - #3

Open
felipeblazing wants to merge 2 commits into
sirius-db:release/26.06from
felipeblazing:perf/groupby-shmem-replication
Open

perf: give each warp lane its own groupby shared-memory aggregation slot#3
felipeblazing wants to merge 2 commits into
sirius-db:release/26.06from
felipeblazing:perf/groupby-shmem-replication

Conversation

@felipeblazing

Copy link
Copy Markdown

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

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
mike-wendt force-pushed the release/26.06 branch 5 times, most recently from eb2e22f to d49c634 Compare August 21, 2026 19:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants