[cudax][cuco] Port fixed_capacity_map benchmarks from cuCollections

[sleeepyjack]

Closes #9595.

Summary

Ports the relevant cuCollections static_map NVBench benchmarks to cudax fixed_capacity_map.

Adjusts the default cooperative-group size for fixed_capacity_map to align with the optimal setup used by cuCollections.

Changes

• Adds shared cudax/cuco benchmark helpers under cudax/benchmarks/bench/cuco/common/.
• Ports static_map insert benchmarks as fixed_capacity_map benchmarks:
  • fixed_capacity_map_insert_unique_capacity
  • fixed_capacity_map_insert_unique_occupancy
  • fixed_capacity_map_insert_uniform_multiplicity
• Ports static_map contains benchmarks as fixed_capacity_map benchmarks:
  • fixed_capacity_map_contains_unique_capacity
  • fixed_capacity_map_contains_unique_occupancy
  • fixed_capacity_map_contains_unique_matching_rate

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds CUCO benchmark defaults and key-generation utilities, then adds NVBench benchmarks for fixed_capacity_map::insert_async and contains_async. The PR also changes the map’s default probing scheme and updates a related shared-memory test.

Changes

fixed_capacity_map benchmark suite

Layer / File(s)	Summary
Benchmark defaults cudax/benchmarks/bench/cuco/common/defaults.cuh	Defines benchmark type axes, scalar defaults, deterministic seed, and input-size, occupancy, multiplicity, and matching-rate sweeps.
Key generator utility cudax/benchmarks/bench/cuco/common/key_generator.cuh	Adds unique and uniform distributions, seeded device-side generation and dropout, NVBench distribution mapping, and type registrations.
insert_async benchmark cudax/benchmarks/bench/cuco/fixed_capacity_map/insert.cu	Generates keys and pairs, constructs a fixed_capacity_map, times insert_async, clears the map afterward, and registers three variants.
contains_async benchmark cudax/benchmarks/bench/cuco/fixed_capacity_map/contains.cu	Builds and populates the map, applies matching-rate dropout, invokes contains_async, and registers three variants.
Default probing scheme change cudax/include/cuda/experimental/__cuco/fixed_capacity_map.cuh, cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu	Changes the default probing scheme to linear_probing<4, hash<_Key>> and updates shared-memory test naming and parameters.

Assessment against linked issues

Objective	Addressed	Explanation
Add benchmarks for fixed_capacity_map::insert and fixed_capacity_map::contains [#9595]	✅

Out-of-scope changes

Code Change	Explanation
Change the default probing scheme to linear_probing<4, hash<_Key>> (cudax/include/cuda/experimental/__cuco/fixed_capacity_map.cuh:73)	The linked objective requests insert and contains benchmarks; it does not specify changing the map’s default probing behavior.

Possibly related PRs

• NVIDIA/cccl#9719 — Changes stream and memory-resource API expectations used by these fixed_capacity_map benchmark paths.

Suggested labels: cudax

Suggested reviewers: caugonnet, srinivasyadav18

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (3)

cudax/benchmarks/bench/cuco/common/key_generator.cuh (1)

207-248: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

suggestion: dropout() unconditionally shuffles [begin, end) at line 247, even when keep_prob >= 1.0 (no keys actually dropped). This side effect isn't mentioned in the docstring (lines 212-223), which only describes the drop/replace behavior. Worth documenting the shuffle's purpose (or gating it behind the keep_prob < 1.0 branch if it's only meant to intersperse dropped keys) so callers relying on this utility understand the full behavior.

As per path instructions for cudax/**/*, review should focus on "experimental API clarity" among other things.

Doc clarification example

   //! Drops keys with probability `1 - keep_prob` using the provided execution policy.
   //!
   //! Replaced keys are sampled from `[N, max_key]`, where `N` is the number of keys in the range.
+  //! The full range is always shuffled afterward, regardless of `keep_prob`.
   //!

Source: Path instructions

cudax/benchmarks/bench/cuco/fixed_capacity_map/insert.cu (1)

46-74: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

suggestion: the device/stream/memory-resource setup, key→pair transform, and map construction (lines 46-74) are duplicated almost verbatim in contains.cu (lines 46-75). Since common/ already houses shared benchmark utilities, consider extracting a shared helper (e.g. a build_map(state, keys) or similar in common/) to avoid drift between the two benchmarks as the map API evolves.

cudax/benchmarks/bench/cuco/common/defaults.cuh (1)

28-39: 🚀 Performance & Scalability | 🔵 Trivial

suggestion: trim the low-occupancy sweep The 10% point drives fixed_capacity_map_insert_unique_occupancy and fixed_capacity_map_contains_unique_occupancy to ~1e9 slots for int64_t/int64_t instances (~16 GiB of slots alone, before overhead), and that multiplies across 9 occupancy values and 2 valid type combinations. Consider raising the minimum occupancy or lowering n so the sweep stays practical on smaller GPUs.

Source: Path instructions

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d0bd7519-6228-406b-9cb5-c72d712b117d

📥 Commits

Reviewing files that changed from the base of the PR and between 3f87270 and 2ea4d34.

📒 Files selected for processing (5)

• cudax/benchmarks/bench/cuco/common/defaults.cuh
• cudax/benchmarks/bench/cuco/common/key_generator.cuh
• cudax/benchmarks/bench/cuco/fixed_capacity_map/contains.cu
• cudax/benchmarks/bench/cuco/fixed_capacity_map/insert.cu
• cudax/include/cuda/experimental/__cuco/fixed_capacity_map.cuh

[PointKernel]

@sleeepyjack can you please share the benchmark results comparing cuco::static_map and cudax::cuco::fixed_capacity_map? Just wanted to make sure we didn't bring any surprises during the migration.

[sleeepyjack]

CI fails due to #9755 . I'm adding a temporary workaround.

[sleeepyjack]

Benchmark/codegen comparison

Compared CUCO static_map upstream/dev against this PR's cudax fixed_capacity_map port on an RTX PRO 6000 Blackwell, CUDA 13.3, GCC 14, Release sm_120.

Δ Time is CCCL vs CUCO NVBench cold mean GPU time; lower is better. Δ Throughput is CCCL vs CUCO element rate; higher is better.

• Compared 56 matched unique-capacity/unique-occupancy measurements.
• Geomean Δ Time: +0.7% over the included measurements.
• Remaining small-input I64/I64 unique-capacity deltas are likely affected by CCCL not yet having CUCO's 128-bit atomics support from NVIDIA/cuCollections#803.

static_map_insert_unique_capacity → fixed_capacity_map_insert_unique_capacity

Geomean Δ Time: +2.1%

Key	NumInputs	Occupancy	CUCO GPU	CCCL GPU	Δ Time	CUCO Elem/s	CCCL Elem/s	Δ Throughput	CUCO Noise	CCCL Noise
I32	8000	0.5	14.254 us	14.370 us	+0.8%	561.237 M/s	556.732 M/s	-0.8%	2.81%	1.82%
I32	80000	0.5	17.050 us	15.428 us	-9.5%	4.692 G/s	5.185 G/s	+10.5%	5.63%	6.62%
I32	800000	0.5	79.862 us	81.930 us	+2.6%	10.017 G/s	9.764 G/s	-2.5%	0.34%	0.17%
I32	8000000	0.5	1.011 ms	1.015 ms	+0.4%	7.911 G/s	7.883 G/s	-0.4%	0.14%	0.12%
I32	80000000	0.5	18.668 ms	18.652 ms	-0.1%	4.285 G/s	4.289 G/s	+0.1%	0.02%	0.07%
I64	8000	0.5	14.296 us	16.263 us	+13.8%	559.595 M/s	491.915 M/s	-12.1%	2.68%	3.22%
I64	80000	0.5	17.240 us	18.440 us	+7.0%	4.640 G/s	4.338 G/s	-6.5%	5.86%	3.33%
I64	800000	0.5	92.323 us	98.304 us	+6.5%	8.665 G/s	8.138 G/s	-6.1%	0.63%	0.18%
I64	8000000	0.5	1.492 ms	1.509 ms	+1.2%	5.363 G/s	5.300 G/s	-1.2%	0.08%	0.09%
I64	80000000	0.5	19.466 ms	19.447 ms	-0.1%	4.110 G/s	4.114 G/s	+0.1%	0.03%	0.03%

static_map_insert_unique_occupancy → fixed_capacity_map_insert_unique_occupancy

Geomean Δ Time: -0.0%

Key	NumInputs	Occupancy	CUCO GPU	CCCL GPU	Δ Time	CUCO Elem/s	CCCL Elem/s	Δ Throughput	CUCO Noise	CCCL Noise
I32	100000000	0.1	23.760 ms	23.796 ms	+0.2%	4.209 G/s	4.202 G/s	-0.2%	0.06%	0.02%
I32	100000000	0.2	23.620 ms	23.627 ms	+0.0%	4.234 G/s	4.232 G/s	-0.0%	0.03%	0.07%
I32	100000000	0.3	23.508 ms	23.526 ms	+0.1%	4.254 G/s	4.251 G/s	-0.1%	0.06%	0.05%
I32	100000000	0.4	23.457 ms	23.477 ms	+0.1%	4.263 G/s	4.259 G/s	-0.1%	0.05%	0.01%
I32	100000000	0.5	23.521 ms	23.513 ms	-0.0%	4.252 G/s	4.253 G/s	+0.0%	0.04%	0.06%
I32	100000000	0.6	23.728 ms	23.739 ms	+0.0%	4.214 G/s	4.212 G/s	-0.0%	0.04%	0.04%
I32	100000000	0.7	24.231 ms	24.258 ms	+0.1%	4.127 G/s	4.122 G/s	-0.1%	0.03%	0.07%
I32	100000000	0.8	25.278 ms	25.277 ms	-0.0%	3.956 G/s	3.956 G/s	+0.0%	0.06%	0.03%
I32	100000000	0.9	27.343 ms	27.458 ms	+0.4%	3.657 G/s	3.642 G/s	-0.4%	0.03%	0.03%
I64	100000000	0.1	24.465 ms	24.418 ms	-0.2%	4.087 G/s	4.095 G/s	+0.2%	0.06%	0.07%
I64	100000000	0.2	24.369 ms	24.321 ms	-0.2%	4.104 G/s	4.112 G/s	+0.2%	0.06%	0.09%
I64	100000000	0.3	24.290 ms	24.268 ms	-0.1%	4.117 G/s	4.121 G/s	+0.1%	0.06%	0.04%
I64	100000000	0.4	24.333 ms	24.272 ms	-0.3%	4.110 G/s	4.120 G/s	+0.3%	0.03%	0.07%
I64	100000000	0.5	24.439 ms	24.430 ms	-0.0%	4.092 G/s	4.093 G/s	+0.0%	0.04%	0.02%
I64	100000000	0.6	24.752 ms	24.721 ms	-0.1%	4.040 G/s	4.045 G/s	+0.1%	0.05%	0.04%
I64	100000000	0.7	25.318 ms	25.279 ms	-0.2%	3.950 G/s	3.956 G/s	+0.2%	0.04%	0.05%
I64	100000000	0.8	26.133 ms	26.129 ms	-0.0%	3.827 G/s	3.827 G/s	+0.0%	0.08%	0.01%
I64	100000000	0.9	28.599 ms	28.629 ms	+0.1%	3.497 G/s	3.493 G/s	-0.1%	0.04%	0.04%

static_map_contains_unique_capacity → fixed_capacity_map_contains_unique_capacity

Geomean Δ Time: +2.0%

Key	NumInputs	Occupancy	CUCO GPU	CCCL GPU	Δ Time	CUCO Elem/s	CCCL Elem/s	Δ Throughput	CUCO Noise	CCCL Noise
I32	8000	0.5	6.236 us	6.162 us	-1.2%	1.283 G/s	1.298 G/s	+1.2%	6.45%	3.14%
I32	80000	0.5	8.210 us	8.192 us	-0.2%	9.744 G/s	9.766 G/s	+0.2%	1.36%	0.10%
I32	800000	0.5	35.099 us	34.866 us	-0.7%	22.793 G/s	22.945 G/s	+0.7%	2.00%	0.91%
I32	8000000	0.5	366.171 us	366.127 us	-0.0%	21.848 G/s	21.850 G/s	+0.0%	0.24%	0.25%
I32	80000000	0.5	5.420 ms	5.419 ms	-0.0%	14.759 G/s	14.763 G/s	+0.0%	0.02%	0.02%
I64	8000	0.5	6.146 us	7.483 us	+21.7%	1.302 G/s	1.069 G/s	-17.9%	0.59%	13.00%
I64	80000	0.5	8.189 us	8.237 us	+0.6%	9.769 G/s	9.712 G/s	-0.6%	1.69%	2.48%
I64	800000	0.5	38.167 us	38.912 us	+2.0%	20.961 G/s	20.559 G/s	-1.9%	2.59%	0.03%
I64	8000000	0.5	483.131 us	483.190 us	+0.0%	16.559 G/s	16.557 G/s	-0.0%	0.14%	0.11%
I64	80000000	0.5	5.454 ms	5.458 ms	+0.1%	14.667 G/s	14.657 G/s	-0.1%	0.05%	0.02%

static_map_contains_unique_occupancy → fixed_capacity_map_contains_unique_occupancy

Geomean Δ Time: +0.1%

Key	NumInputs	Occupancy	CUCO GPU	CCCL GPU	Δ Time	CUCO Elem/s	CCCL Elem/s	Δ Throughput	CUCO Noise	CCCL Noise
I32	100000000	0.1	6.648 ms	6.656 ms	+0.1%	15.042 G/s	15.024 G/s	-0.1%	0.02%	0.02%
I32	100000000	0.2	6.625 ms	6.631 ms	+0.1%	15.095 G/s	15.081 G/s	-0.1%	0.02%	0.02%
I32	100000000	0.3	6.627 ms	6.633 ms	+0.1%	15.089 G/s	15.076 G/s	-0.1%	0.02%	0.03%
I32	100000000	0.4	6.676 ms	6.680 ms	+0.1%	14.979 G/s	14.970 G/s	-0.1%	0.02%	0.02%
I32	100000000	0.5	6.801 ms	6.809 ms	+0.1%	14.705 G/s	14.686 G/s	-0.1%	0.02%	0.02%
I32	100000000	0.6	7.064 ms	7.065 ms	+0.0%	14.157 G/s	14.154 G/s	-0.0%	0.02%	0.02%
I32	100000000	0.7	7.582 ms	7.583 ms	+0.0%	13.190 G/s	13.187 G/s	-0.0%	0.01%	0.02%
I32	100000000	0.8	8.734 ms	8.739 ms	+0.1%	11.450 G/s	11.443 G/s	-0.1%	0.01%	0.02%
I32	100000000	0.9	12.357 ms	12.398 ms	+0.3%	8.092 G/s	8.066 G/s	-0.3%	0.01%	0.02%
I64	100000000	0.1	6.659 ms	6.662 ms	+0.0%	15.017 G/s	15.010 G/s	-0.0%	0.06%	0.03%
I64	100000000	0.2	6.639 ms	6.638 ms	-0.0%	15.062 G/s	15.064 G/s	+0.0%	0.05%	0.03%
I64	100000000	0.3	6.647 ms	6.646 ms	-0.0%	15.045 G/s	15.046 G/s	+0.0%	0.02%	0.08%
I64	100000000	0.4	6.701 ms	6.702 ms	+0.0%	14.924 G/s	14.920 G/s	-0.0%	0.03%	0.03%
I64	100000000	0.5	6.839 ms	6.838 ms	-0.0%	14.623 G/s	14.625 G/s	+0.0%	0.03%	0.03%
I64	100000000	0.6	7.099 ms	7.100 ms	+0.0%	14.086 G/s	14.085 G/s	-0.0%	0.02%	0.02%
I64	100000000	0.7	7.633 ms	7.630 ms	-0.0%	13.101 G/s	13.106 G/s	+0.0%	0.02%	0.01%
I64	100000000	0.8	8.822 ms	8.819 ms	-0.0%	11.335 G/s	11.340 G/s	+0.0%	0.01%	0.01%
I64	100000000	0.9	12.663 ms	12.671 ms	+0.1%	7.897 G/s	7.892 G/s	-0.1%	0.02%	0.01%

Codegen

• Timed I64 insert kernel register usage drops from 32 registers in CUCO to 30 in the cudax port (sm_120); timed contains stays at 30 registers.
• Timed I32 contains register usage is unchanged

[PointKernel]

Two small questions but not blocking.

[sleeepyjack]

Since the benchmarks are considered user-facing and not directly part of the library, should I switch to using cuda:: and thrust:: directly instead of forcing top-level scope through ::cuda::?

[github-actions]

😬 CI Workflow Results

🟥 Finished in 1h 09m: Pass: 98%/57 | Total: 15h 43m | Max: 1h 09m | Hits: 34%/105613

See results here.