[BUG]: static_multimap::retrieve gives wrong output with custom equality/hasher

[fkallen]

Is this a duplicate?

• I confirmed there appear to be no duplicate issues for this bug (https://github.com/NVIDIA/cuCollections/issues)

Type of Bug

Silent Failure

Describe the bug

@nmetzinger and I were debugging some piece of code and found the issue to be wrong results retrieved from a static_multimap.

We are using a custom hasher / equality operator for retrieve to associate additional data with the retrieved values via the output_probe buffer. However, retrieve results are only partially correct.

Below is the reproducer. Hash table contents are loaded from file. The file is included in the attached self-contained example. With the current cuCollections commit, CUDA 13.3 with its included CCCL, we get the following output on sm_120.

loaded 7909 pairs
d_numHitsPerKey:
11 0 0 0 0 0 0 0 0 0 6 3 0 
numReads 1
numKeys 13
numHits 20

with original keys
d_keysOfHits:
23846566 23846566 23846566 23846566 23846566 23846566 24892615 24892615 24892615 8083701 8083701 8083701 8083701 8083701 8083701 8083701 8083701 8083701 8083701 8083701 
d_hits:
4412 6106 3995 4756 7175 2621 4443 3996 7176 1426 2821 1669 3370 6454 6592 4039 5454 6546 3541 4345 


with custom keys
querying h_customKeys:
(8083701,0) (23797427,1) (23320168,2) (8047878,3) (22651094,4) (20191956,5) (8608406,6) (7033559,7) (23747319,8) (21716725,9) (23846566,10) (24892615,11) (24811765,12) 

d_customKeysOfHits:
(23846566,10) (23846566,10) (23846566,10) (23846566,10) (23846566,10) (23846566,10) (24892615,11) (24892615,11) (24892615,11) (8083701,0) (8083701,0) (8083701,0) (8083701,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) 
d_hits:
4412 6106 3995 4756 7175 2621 4443 3996 7176 1426 2821 1669 3370 0 0 0 0 0 0 0 

Note that when using the custom approach, only 4 out of 11 values for key 8083701 are retrieved correctly.

// nvcc -std=c++17 -O0 -g -lineinfo --extended-lambda --expt-relaxed-constexpr -gencode=arch=compute_120,code=sm_120 -I cuCollections/include/ reproducer.cu -o reproducer
#include <cooperative_groups.h>
#include <cuco/static_multimap.cuh>
#include <fstream>
#include <iostream>
#include <limits>
#include <thrust/copy.h>
#include <thrust/device_vector.h>
#include <thrust/iterator/zip_iterator.h>
#include <vector>

using key_type = uint64_t;

struct Location {
    Location() = default;

    __host__ __device__ constexpr Location(uint64_t data_)
: data(data_){}

    uint64_t data;
};

std::ostream& operator<<(std::ostream& os, const Location& loc) {
    os << loc.data;
    return os;
}

constexpr key_type empty_key_sentinel = std::numeric_limits<key_type>::max();
constexpr Location empty_value_sentinel{ std::numeric_limits<uint64_t>::max()};

using HashTable = cuco::static_multimap<key_type, Location>;

using CustomRetrieveKeyType = cuda::std::tuple<key_type, int>;


struct CustomKeyEqual {
    __host__ __device__ bool operator()(const CustomRetrieveKeyType& customKey,
                                        const key_type& key) const {
        return cuda::std::get<0>(customKey) == key;
    }

};

template<int hasherNum>
struct CustomKeyHash {

    using Hasher = typename HashTable::hasher;

    __host__ __device__ auto operator()(const CustomRetrieveKeyType& customKey) const {
        auto key = cuda::std::get<0>(customKey);

        auto hasher = cuda::std::get<hasherNum>(Hasher{});
        return hasher(key);
    }
};



template <typename MapRef, typename InputIterator, typename OutputIterator>
__global__ void getNumHitsPerKeyKernel(MapRef map, InputIterator keys,
                                       std::size_t n,
                                       OutputIterator numHitsPerKey) {
    namespace cg = cooperative_groups;

    constexpr auto cg_size = MapRef::cg_size;

    auto tile =
        cg::tiled_partition<cg_size, cg::thread_block>(cg::this_thread_block());

    int64_t const loop_stride = gridDim.x * blockDim.x / cg_size;
    int64_t idx = (blockDim.x * blockIdx.x + threadIdx.x) / cg_size;

    while (idx < n) {
        int keycount = map.count(tile, *(keys + idx));
        keycount = cg::reduce(tile, keycount, cg::plus<int>{});
        if (tile.thread_rank() == 0) {
            // uint64_t key = *(keys + idx);
            // printf("idx %lu, key %lu, count %d\n", idx, key, keycount);
            numHitsPerKey[idx] = keycount;
        }
        idx += loop_stride;
    }
}


template <typename InputIterator, typename OutputIterator>
void callGetNumHitsPerKeyKernel(HashTable& map, InputIterator d_keys,
                                std::size_t numKeys,
                                OutputIterator d_numHitsPerKey,
                                cudaStream_t stream) {
    auto map_ref_count = map.ref(cuco::count);
    int cg_size = decltype(map_ref_count)::cg_size;
    int blocksize = 128;
    int gridsize = cuda::ceil_div(numKeys, blocksize / cg_size);
    getNumHitsPerKeyKernel<<<gridsize, blocksize, 0, stream>>>(
        map_ref_count, d_keys, numKeys, d_numHitsPerKey);
}


int main() {
    std::string filename = "reference_kvpairs.bin";

    using KVPair = cuco::pair<key_type, Location>;

    std::vector<KVPair> reference_kvpairs;
    std::ifstream infile(filename);
    size_t numrefpairs;
    infile.read((char*)&numrefpairs, sizeof(size_t));
    reference_kvpairs.resize(numrefpairs);
    infile.read((char*)reference_kvpairs.data(), sizeof(KVPair) * numrefpairs);

    std::cout << "loaded " << numrefpairs << " pairs\n";

    double loadFactor = 0.5;
    HashTable gpuHashTable(numrefpairs / loadFactor,
                           cuco::empty_key{empty_key_sentinel},
                           cuco::empty_value{empty_value_sentinel});

    thrust::device_vector<KVPair> d_kvpairs = reference_kvpairs;

    gpuHashTable.insert(d_kvpairs.begin(), d_kvpairs.end());

    // querying
    thrust::device_vector<key_type> d_queryKeys{
        8083701ull,  23797427ull, 23320168ull, 8047878ull,  22651094ull,
        20191956ull, 8608406ull,  7033559ull,  23747319ull, 21716725ull,
        23846566ull, 24892615ull, 24811765ull};
    thrust::device_vector<int> d_queryOtherData{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
    assert(d_queryKeys.size() == d_queryOtherData.size());

    // thrust::reverse(d_queryKeys.begin(), d_queryKeys.end());
    // thrust::reverse(d_queryOtherData.begin(), d_queryOtherData.end());

    const int numReads = 1;
    const int numKeys = d_queryKeys.size();
    thrust::device_vector<int> d_numHitsPerKey(numKeys);

    callGetNumHitsPerKeyKernel(gpuHashTable, d_queryKeys.begin(), numKeys,
                               d_numHitsPerKey.begin(), (cudaStream_t)0);

    std::cout << "d_numHitsPerKey:\n";
    for (int i = 0; i < numKeys; i++) {
        std::cout << d_numHitsPerKey[i] << " ";
    }
    std::cout << "\n";

    const int numHits =
        thrust::reduce(d_numHitsPerKey.begin(), d_numHitsPerKey.end());

    std::cout << "numReads " << numReads << "\n";
    std::cout << "numKeys " << numKeys << "\n";
    std::cout << "numHits " << numHits << "\n";

    std::cout << "\n";

    std::cout << "with original keys\n";
    {
        thrust::device_vector<Location> d_hits(numHits);
        thrust::device_vector<key_type> d_keysOfHits(numHits);

        gpuHashTable.retrieve(
            d_queryKeys.begin(),
            d_queryKeys.end(),
            d_keysOfHits.begin(),
            cuda::make_transform_output_iterator(
                d_hits.begin(),
                [] __device__(auto cucopair) -> Location { return cucopair.second; }
            )
        );

        cudaDeviceSynchronize();

        std::cout << "d_keysOfHits:\n";
        for (int i = 0; i < numHits; i++) {
            std::cout << d_keysOfHits[i] << " ";
        }
        std::cout << "\n";

        std::cout << "d_hits:\n";
        for (int i = 0; i < numHits; i++) {
            std::cout << d_hits[i] << " ";
        }
        std::cout << "\n";
    }

    std::cout << "\n";
    std::cout << "\n";


    std::cout << "with custom keys\n";

    {
        auto customKeys = thrust::make_zip_iterator(d_queryKeys.begin(),
                                                    d_queryOtherData.begin());

        std::vector<CustomRetrieveKeyType> h_customKeys(numKeys);
        thrust::copy(customKeys, customKeys + numKeys, h_customKeys.begin());

        std::cout << "querying h_customKeys:\n";
        for (int i = 0; i < numKeys; i++) {
            CustomRetrieveKeyType tup = h_customKeys[i];
            std::cout << "(" << cuda::std::get<0>(tup) << ","
                      << cuda::std::get<1>(tup) << ") ";
        }
        std::cout << "\n";
        std::cout << "\n";

        thrust::device_vector<Location> d_hits(numHits);
        thrust::device_vector<CustomRetrieveKeyType>
            d_customKeysOfHits(numHits);

        CustomKeyEqual probe_eq;
        CustomKeyHash<0> probe_hash_0;
        CustomKeyHash<1> probe_hash_1;
        auto probe_doublehasher = cuda::std::make_tuple(probe_hash_0, probe_hash_1);
        gpuHashTable.retrieve(customKeys,
            customKeys + numKeys,
            probe_eq,
            probe_doublehasher,
            d_customKeysOfHits.begin(),
            cuda::make_transform_output_iterator(
                d_hits.begin(),
                [] __device__(auto cucopair) -> Location { return cucopair.second; }
            )
        );

        cudaDeviceSynchronize();

        std::cout << "d_customKeysOfHits:\n";
        for (int i = 0; i < numHits; i++) {
            CustomRetrieveKeyType tup = d_customKeysOfHits[i];
            std::cout << "(" << cuda::std::get<0>(tup) << ","
                      << cuda::std::get<1>(tup) << ") ";
        }
        std::cout << "\n";

        std::cout << "d_hits:\n";
        for (int i = 0; i < numHits; i++) {
            std::cout << d_hits[i] << " ";
        }
        std::cout << "\n";
    }

    std::cout << "\n";
    std::cout << "\n";



}

How to Reproduce

cuco_reproducer.zip

Expected behavior

All results should be correct

Reproduction link

No response

Operating System

Linux

nvidia-smi output

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 610.43.02 KMD Version: 610.43.02 CUDA UMD Version: 13.3 |
+-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA RTX PRO 6000 Blac... On | 00000000:C1:00.0 Off | 0 |
| N/A 29C P8 34W / 600W | 0MiB / 97887MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+------------------------+----------------------+

NVCC version

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2026 NVIDIA Corporation
Built on Fri_Apr_24_07:22:02_PM_PDT_2026
Cuda compilation tools, release 13.3, V13.3.33
Build cuda_13.3.r13.3/compiler.37862127_0

[sleeepyjack]

@fkallen @nmetzinger thanks for filing this issue.

After some investigation I think I found the problem. It's not directly a bug but a quirk in our interface.

The root cause is a hash-seed mismatch between how the map probes and how the reproducer reconstructs the hashers.
The default probing scheme is cuco::double_hashing<8, default_hash_function<Key>>, and its constructor seeds the two hashers asymmetrically to avoid secondary collisions: double_hashing(Hash1 const& hash1 = {}, Hash2 const& hash2 = {1})

So the map actually probes with { xxhash_32{seed=0}, xxhash_32{seed=1} }. But the reproducer rebuilds them via cuda::std::get<hasherNum>(Hasher{}), and Hasher{} default-constructs both tuple elements, giving { xxhash_32{0}, xxhash_32{0}}. The second hasher ends up with seed 0 instead of 1.

TL;DR: the problem boils down to HashTable::hasher{} != gpuHashTable.hash_function(). The default-constructed hasher seeds both sub-hashers with 0, while the map's real hashers are seeded {0, 1}. @PointKernel We should discuss how to defuse this footgun.

The fix is to pull the map's actual (correctly seeded) hashers from hash_function() instead of default-constructing them, and wrap them so they hash the embedded key:

struct CustomKeyHash {
  cuco::default_hash_function<key_type> hasher;  // carries the real seed
  __host__ __device__ auto operator()(CustomRetrieveKeyType const& k) const {
    return hasher(cuda::std::get<0>(k));
  }
};

auto hashers = gpuHashTable.hash_function(); // tuple{ H{0}, H{1} }
CustomKeyHash probe_hash_0{cuda::std::get<0>(hashers)};
CustomKeyHash probe_hash_1{cuda::std::get<1>(hashers)};
auto probe_doublehasher = cuda::std::make_tuple(probe_hash_0, probe_hash_1);

[fkallen]

Thanks @sleeepyjack ! I can confirm the reproducer works with your proposed fixed.

For a user it might be more convenient to not have to deal with the internal hash functions at all, but rather tell the map what the input to the hash function should be. For example, have some kind of transformOp(probekey)->key that can be customized by the user and which is the identity function for standard retrieve. Probing would internally compute defaulthashfunc( transformOp(probeKey)) .

One more question: Do we have to be concerned about the compiler warning that comes with our custom approach or is it safe to ignore it?

cuCollections/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh(1269): warning #20054-D: dynamic initialization is not supported for a function-scope static __shared__ variable within a __device__/__global__ function
      __attribute__((shared)) cuco::pair<probe_type, value_type> buffers[num_flushing_tiles][buffer_size];

[PointKernel]

Oh my, thanks @sleeepyjack for catching this. The retrieve kernel is extensively tested in our production code, particularly in cudf join operations, so it's less likely that the bug is in the kernel itself. However, this is a subtle issue that can be easy to overlook.

cuCollections/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh(1269): warning #20054-D: dynamic initialization is not supported for a function-scope static __shared__ variable within a __device__/__global__ function
      __attribute__((shared)) cuco::pair<probe_type, value_type> buffers[num_flushing_tiles][buffer_size];

@fkallen this woudn't be an issue for your use case but let me put up a PR to fix it.

[fkallen]

We got around to apply the fix to our program and noticed some bogus behavior when using the device-side retrieve API in a custom kernel. Depending on whether we printf before the retrieve (with optionally syncing the block) , the retrieved results change. Maybe you can comment on this behavior, it is not clear to us if it is again an error on our side or if it is a compiler issue or cuCollections issue.

with custom retrieve kernel
d_customKeysOfHits:
(23846566,10) (23846566,10) (23846566,10) (23846566,10) (23846566,10) (23846566,10) (24892615,11) (24892615,11) (24892615,11) (8083701,0) (8083701,0) (8083701,0) (8083701,0) (8083701,0) (8083701,0) (8083701,0) (8083701,0) (8083701,0) (8083701,0) (8083701,0) 
d_hits:
4412 6106 2621 4756 7175 3995 4443 7176 3996 3541 5454 4039 4345 6454 6546 1669 3370 1426 6592 2821 


with custom retrieve kernel, thread 0 prints before retrieve func
in kernel. idx 0, querying keys:
(8083701, 0) (23797427, 1) (23320168, 2) (8047878, 3) (22651094, 4) (20191956, 5) (8608406, 6) (7033559, 7) (23747319, 8) (21716725, 9) (23846566, 10) (24892615, 11) (24811765, 12) 
d_customKeysOfHits:
(23846566,10) (23846566,10) (23846566,10) (23846566,10) (23846566,10) (23846566,10) (24892615,11) (24892615,11) (24892615,11) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) 
d_hits:
4412 6106 2621 4756 7175 3995 4443 7176 3996 0 0 0 0 0 0 0 0 0 0 0 


with custom retrieve kernel, thread 0 prints and block synchronizes before retrieve func
in kernel. idx 0, querying keys:
(8083701, 0) (23797427, 1) (23320168, 2) (8047878, 3) (22651094, 4) (20191956, 5) (8608406, 6) (7033559, 7) (23747319, 8) (21716725, 9) (23846566, 10) (24892615, 11) (24811765, 12) 
d_customKeysOfHits:
(0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) (0,0) 
d_hits:
0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0

The reproducer kernel:

template <int bugtype, int block_size, typename MapRef, typename InputIterator,
          typename KeyOffsetIterator, typename LocationOffsetIterator,
          typename OutputIterator, typename OutputKeyIterator>
__global__ void retrieveHitsCustomKernel(MapRef map, InputIterator keys,
                                   KeyOffsetIterator numKeysPerReadPS,
                                   LocationOffsetIterator numHitsPerReadPS,
                                   int numReads, OutputIterator locations,
                                   OutputKeyIterator keysOfLocations) {
    namespace cg = cooperative_groups;

    auto cg_block = cg::this_thread_block();

    int64_t const loop_stride = gridDim.x;
    int64_t idx = blockIdx.x;

    __shared__ cuda::atomic<int> atomic_counter;

    while (idx < numReads) {

        __syncthreads();
        if (threadIdx.x == 0) {
            atomic_counter = 0;
        }
        __syncthreads();

        auto key_start = numKeysPerReadPS[idx];
        auto key_end = numKeysPerReadPS[idx + 1];

        if constexpr (bugtype > 0){
            if (threadIdx.x == 0) {
                printf("in kernel. idx %lu, querying keys:\n", idx);
                for (auto i = key_start; i != key_end; i++) {
                    auto tuple = keys[i];
                    printf("(%lu, %d) ", size_t(cuda::std::get<0>(tuple)),
                        int(cuda::std::get<1>(tuple)));
                }
                printf("\n");
            }
        }
        if constexpr (bugtype > 1){
            __syncthreads();
        }

        map.retrieve<block_size>(
            cg_block, keys + key_start, keys + key_end, keysOfLocations + numHitsPerReadPS[idx],
            locations + numHitsPerReadPS[idx], atomic_counter);

        idx += loop_stride;
    }
}

template <int bugtype, typename InputIterator, typename CustomKeyComparator,
          typename CustomHasher, typename KeyOffsetIterator,
          typename LocationOffsetIterator, typename OutputIterator,
          typename OutputKeyIterator>
void callRetrieveHitsKernel(HashTable& map, InputIterator d_keys,
                            size_t numKeys, CustomKeyComparator probe_equal,
                            CustomHasher probe_hash,
                            KeyOffsetIterator d_numKeysPerReadPS,
                            LocationOffsetIterator d_numHitsPerReadPS,
                            int numReads, OutputIterator d_locations,
                            OutputKeyIterator d_keysOfLocations,
                            cudaStream_t stream) {
    auto map_ref_retrieve = map.ref(cuco::retrieve)
                                .rebind_key_eq(probe_equal)
                                .rebind_hash_function(probe_hash);

    constexpr int blocksize1 = 128;
    int gridsize1 = cuda::ceil_div(numReads, blocksize1);

    retrieveHitsCustomKernel<bugtype, blocksize1><<<gridsize1, blocksize1, 0, stream>>>(
        map_ref_retrieve, d_keys, d_numKeysPerReadPS, d_numHitsPerReadPS,
        numReads, d_locations, d_keysOfLocations);
}

The complete modified reproducer is below. (uses the same hash table data as before)

// nvcc -std=c++17 -O0 -g -lineinfo --extended-lambda --expt-relaxed-constexpr -gencode=arch=compute_120,code=sm_120 -I cuCollections/include/ reproducer.cu -o reproducer
#include <cooperative_groups.h>
#include <cuco/static_multimap.cuh>
#include <fstream>
#include <iostream>
#include <limits>
#include <thrust/copy.h>
#include <thrust/device_vector.h>
#include <thrust/iterator/zip_iterator.h>
#include <vector>

using key_type = uint64_t;

struct Location {
    Location() = default;

    __host__ __device__ constexpr Location(uint64_t data_)
: data(data_){}

    uint64_t data;
};

std::ostream& operator<<(std::ostream& os, const Location& loc) {
    os << loc.data;
    return os;
}

constexpr key_type empty_key_sentinel = std::numeric_limits<key_type>::max();
constexpr Location empty_value_sentinel{ std::numeric_limits<uint64_t>::max()};

using HashTable = cuco::static_multimap<key_type, Location>;

using CustomRetrieveKeyType = cuda::std::tuple<key_type, int>;


struct CustomKeyEqual {
    __host__ __device__ bool operator()(const CustomRetrieveKeyType& customKey,
                                        const key_type& key) const {
        return cuda::std::get<0>(customKey) == key;
    }

};

template<class HashFunc>
struct CustomKeyHash {

    HashFunc hasher;

    CustomKeyHash(const HashFunc& hasher_) : hasher(hasher_){}

    __host__ __device__ auto operator()(const CustomRetrieveKeyType& customKey) const {
        auto key = cuda::std::get<0>(customKey);

        return hasher(key);
    }
};




template <typename MapRef, typename InputIterator, typename OutputIterator>
__global__ void getNumHitsPerKeyKernel(MapRef map, InputIterator keys,
                                       std::size_t n,
                                       OutputIterator numHitsPerKey) {
    namespace cg = cooperative_groups;

    constexpr auto cg_size = MapRef::cg_size;

    auto tile =
        cg::tiled_partition<cg_size, cg::thread_block>(cg::this_thread_block());

    int64_t const loop_stride = gridDim.x * blockDim.x / cg_size;
    int64_t idx = (blockDim.x * blockIdx.x + threadIdx.x) / cg_size;

    while (idx < n) {
        int keycount = map.count(tile, *(keys + idx));
        keycount = cg::reduce(tile, keycount, cg::plus<int>{});
        if (tile.thread_rank() == 0) {
            // uint64_t key = *(keys + idx);
            // printf("idx %lu, key %lu, count %d\n", idx, key, keycount);
            numHitsPerKey[idx] = keycount;
        }
        idx += loop_stride;
    }
}


template <typename InputIterator, typename OutputIterator>
void callGetNumHitsPerKeyKernel(HashTable& map, InputIterator d_keys,
                                std::size_t numKeys,
                                OutputIterator d_numHitsPerKey,
                                cudaStream_t stream) {
    auto map_ref_count = map.ref(cuco::count);
    int cg_size = decltype(map_ref_count)::cg_size;
    int blocksize = 128;
    int gridsize = cuda::ceil_div(numKeys, blocksize / cg_size);
    getNumHitsPerKeyKernel<<<gridsize, blocksize, 0, stream>>>(
        map_ref_count, d_keys, numKeys, d_numHitsPerKey);
}


template <int bugtype, int block_size, typename MapRef, typename InputIterator,
          typename KeyOffsetIterator, typename LocationOffsetIterator,
          typename OutputIterator, typename OutputKeyIterator>
__global__ void retrieveHitsCustomKernel(MapRef map, InputIterator keys,
                                   KeyOffsetIterator numKeysPerReadPS,
                                   LocationOffsetIterator numHitsPerReadPS,
                                   int numReads, OutputIterator locations,
                                   OutputKeyIterator keysOfLocations) {
    namespace cg = cooperative_groups;

    auto cg_block = cg::this_thread_block();

    int64_t const loop_stride = gridDim.x;
    int64_t idx = blockIdx.x;

    __shared__ cuda::atomic<int> atomic_counter;

    while (idx < numReads) {

        __syncthreads();
        if (threadIdx.x == 0) {
            atomic_counter = 0;
        }
        __syncthreads();

        auto key_start = numKeysPerReadPS[idx];
        auto key_end = numKeysPerReadPS[idx + 1];

        if constexpr (bugtype > 0){
            if (threadIdx.x == 0) {
                printf("in kernel. idx %lu, querying keys:\n", idx);
                for (auto i = key_start; i != key_end; i++) {
                    auto tuple = keys[i];
                    printf("(%lu, %d) ", size_t(cuda::std::get<0>(tuple)),
                        int(cuda::std::get<1>(tuple)));
                }
                printf("\n");
            }
        }
        if constexpr (bugtype > 1){
            __syncthreads();
        }

        map.retrieve<block_size>(
            cg_block, keys + key_start, keys + key_end, keysOfLocations + numHitsPerReadPS[idx],
            locations + numHitsPerReadPS[idx], atomic_counter);

        idx += loop_stride;
    }
}

template <int bugtype, typename InputIterator, typename CustomKeyComparator,
          typename CustomHasher, typename KeyOffsetIterator,
          typename LocationOffsetIterator, typename OutputIterator,
          typename OutputKeyIterator>
void callRetrieveHitsKernel(HashTable& map, InputIterator d_keys,
                            size_t numKeys, CustomKeyComparator probe_equal,
                            CustomHasher probe_hash,
                            KeyOffsetIterator d_numKeysPerReadPS,
                            LocationOffsetIterator d_numHitsPerReadPS,
                            int numReads, OutputIterator d_locations,
                            OutputKeyIterator d_keysOfLocations,
                            cudaStream_t stream) {
    auto map_ref_retrieve = map.ref(cuco::retrieve)
                                .rebind_key_eq(probe_equal)
                                .rebind_hash_function(probe_hash);

    constexpr int blocksize1 = 128;
    int gridsize1 = cuda::ceil_div(numReads, blocksize1);

    retrieveHitsCustomKernel<bugtype, blocksize1><<<gridsize1, blocksize1, 0, stream>>>(
        map_ref_retrieve, d_keys, d_numKeysPerReadPS, d_numHitsPerReadPS,
        numReads, d_locations, d_keysOfLocations);
}


int main() {
    std::string filename = "reference_kvpairs.bin";

    using KVPair = cuco::pair<key_type, Location>;

    std::vector<KVPair> reference_kvpairs;
    std::ifstream infile(filename);
    size_t numrefpairs;
    infile.read((char*)&numrefpairs, sizeof(size_t));
    reference_kvpairs.resize(numrefpairs);
    infile.read((char*)reference_kvpairs.data(), sizeof(KVPair) * numrefpairs);

    std::cout << "loaded " << numrefpairs << " pairs\n";

    double loadFactor = 0.5;
    HashTable gpuHashTable(numrefpairs / loadFactor,
                           cuco::empty_key{empty_key_sentinel},
                           cuco::empty_value{empty_value_sentinel});

    thrust::device_vector<KVPair> d_kvpairs = reference_kvpairs;

    gpuHashTable.insert(d_kvpairs.begin(), d_kvpairs.end());

    // querying
    thrust::device_vector<key_type> d_queryKeys{
        8083701ull,  23797427ull, 23320168ull, 8047878ull,  22651094ull,
        20191956ull, 8608406ull,  7033559ull,  23747319ull, 21716725ull,
        23846566ull, 24892615ull, 24811765ull};
    thrust::device_vector<int> d_queryOtherData{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
    assert(d_queryKeys.size() == d_queryOtherData.size());

    // thrust::reverse(d_queryKeys.begin(), d_queryKeys.end());
    // thrust::reverse(d_queryOtherData.begin(), d_queryOtherData.end());

    const int numReads = 1;
    const int numKeys = d_queryKeys.size();
    thrust::device_vector<int> d_numHitsPerKey(numKeys);

    callGetNumHitsPerKeyKernel(gpuHashTable, d_queryKeys.begin(), numKeys,
                               d_numHitsPerKey.begin(), (cudaStream_t)0);

    std::cout << "d_numHitsPerKey:\n";
    for (int i = 0; i < numKeys; i++) {
        std::cout << d_numHitsPerKey[i] << " ";
    }
    std::cout << "\n";

    const int numHits =
        thrust::reduce(d_numHitsPerKey.begin(), d_numHitsPerKey.end());

    std::cout << "numReads " << numReads << "\n";
    std::cout << "numKeys " << numKeys << "\n";
    std::cout << "numHits " << numHits << "\n";

    std::cout << "\n";

    std::cout << "with original keys\n";
    {
        thrust::device_vector<Location> d_hits(numHits);
        thrust::device_vector<key_type> d_keysOfHits(numHits);

        gpuHashTable.retrieve(
            d_queryKeys.begin(),
            d_queryKeys.end(),
            d_keysOfHits.begin(),
            cuda::make_transform_output_iterator(
                d_hits.begin(),
                [] __device__(auto cucopair) -> Location { return cucopair.second; }
            )
        );

        cudaDeviceSynchronize();

        std::cout << "d_keysOfHits:\n";
        for (int i = 0; i < numHits; i++) {
            std::cout << d_keysOfHits[i] << " ";
        }
        std::cout << "\n";

        std::cout << "d_hits:\n";
        for (int i = 0; i < numHits; i++) {
            std::cout << d_hits[i] << " ";
        }
        std::cout << "\n";
    }

    std::cout << "\n";
    std::cout << "\n";


    std::cout << "with custom keys\n";

    {
        auto customKeys = thrust::make_zip_iterator(d_queryKeys.begin(),
                                                    d_queryOtherData.begin());

        std::vector<CustomRetrieveKeyType> h_customKeys(numKeys);
        thrust::copy(customKeys, customKeys + numKeys, h_customKeys.begin());

        std::cout << "querying h_customKeys:\n";
        for (int i = 0; i < numKeys; i++) {
            CustomRetrieveKeyType tup = h_customKeys[i];
            std::cout << "(" << cuda::std::get<0>(tup) << ","
                      << cuda::std::get<1>(tup) << ") ";
        }
        std::cout << "\n";
        std::cout << "\n";

        
        CustomKeyEqual probe_eq;
        auto hashers = gpuHashTable.hash_function(); // tuple{ H{0}, H{1} }

        //make sure not to store a (host) reference to hash function which would be used on the device
        using HashFunc0 = typename std::remove_reference_t<decltype(cuda::std::get<0>(hashers))>;
        using HashFunc1 = typename std::remove_reference_t<decltype(cuda::std::get<1>(hashers))>;
        CustomKeyHash<HashFunc0> probe_hash_0(cuda::std::get<0>(hashers));
        CustomKeyHash<HashFunc1> probe_hash_1(cuda::std::get<1>(hashers));
        auto probe_doublehasher = cuda::std::make_tuple(probe_hash_0, probe_hash_1);

        {
            std::cout << "with standard retrieve kernel\n";
            thrust::device_vector<Location> d_hits(numHits);
            thrust::device_vector<CustomRetrieveKeyType> d_customKeysOfHits(numHits);


            gpuHashTable.retrieve(customKeys,
                customKeys + numKeys,
                probe_eq,
                probe_doublehasher,
                d_customKeysOfHits.begin(),
                cuda::make_transform_output_iterator(
                    d_hits.begin(),
                    [] __device__(auto cucopair) -> Location { return cucopair.second; }
                )
            );

            cudaDeviceSynchronize();

            std::cout << "d_customKeysOfHits:\n";
            for (int i = 0; i < numHits; i++) {
                CustomRetrieveKeyType tup = d_customKeysOfHits[i];
                std::cout << "(" << cuda::std::get<0>(tup) << ","
                        << cuda::std::get<1>(tup) << ") ";
            }
            std::cout << "\n";

            std::cout << "d_hits:\n";
            for (int i = 0; i < numHits; i++) {
                std::cout << d_hits[i] << " ";
            }
            std::cout << "\n";
        }

        {
            std::cout << "with custom retrieve kernel\n";
            thrust::device_vector<Location> d_hits(numHits);
            thrust::device_vector<CustomRetrieveKeyType> d_customKeysOfHits(numHits);

            thrust::device_vector<int> d_numKeysPerReadPrefixSum({0, numKeys});
            thrust::device_vector<int> d_numHitsPerReadPrefixSum({0, numHits});

            constexpr int bugtype = 0;
            callRetrieveHitsKernel<bugtype>(gpuHashTable, customKeys,
                numKeys, probe_eq,
                probe_doublehasher,
                d_numKeysPerReadPrefixSum.begin(),
                d_numHitsPerReadPrefixSum.begin(),
                numReads, 
                cuda::make_transform_output_iterator(
                    d_hits.begin(),
                    [] __device__(auto cucopair) -> Location { return cucopair.second; }
                ),
                d_customKeysOfHits.begin(),
                (cudaStream_t)0
            );

            cudaDeviceSynchronize();

            std::cout << "d_customKeysOfHits:\n";
            for (int i = 0; i < numHits; i++) {
                CustomRetrieveKeyType tup = d_customKeysOfHits[i];
                std::cout << "(" << cuda::std::get<0>(tup) << ","
                        << cuda::std::get<1>(tup) << ") ";
            }
            std::cout << "\n";

            std::cout << "d_hits:\n";
            for (int i = 0; i < numHits; i++) {
                std::cout << d_hits[i] << " ";
            }
            std::cout << "\n";

            std::cout << "\n";
            std::cout << "\n";
        }

        {
            std::cout << "with custom retrieve kernel, thread 0 prints before retrieve func\n";
            thrust::device_vector<Location> d_hits(numHits);
            thrust::device_vector<CustomRetrieveKeyType> d_customKeysOfHits(numHits);

            thrust::device_vector<int> d_numKeysPerReadPrefixSum({0, numKeys});
            thrust::device_vector<int> d_numHitsPerReadPrefixSum({0, numHits});

            constexpr int bugtype = 1;
            callRetrieveHitsKernel<bugtype>(gpuHashTable, customKeys,
                numKeys, probe_eq,
                probe_doublehasher,
                d_numKeysPerReadPrefixSum.begin(),
                d_numHitsPerReadPrefixSum.begin(),
                numReads, 
                cuda::make_transform_output_iterator(
                    d_hits.begin(),
                    [] __device__(auto cucopair) -> Location { return cucopair.second; }
                ),
                d_customKeysOfHits.begin(),
                (cudaStream_t)0
            );

            cudaDeviceSynchronize();

            std::cout << "d_customKeysOfHits:\n";
            for (int i = 0; i < numHits; i++) {
                CustomRetrieveKeyType tup = d_customKeysOfHits[i];
                std::cout << "(" << cuda::std::get<0>(tup) << ","
                        << cuda::std::get<1>(tup) << ") ";
            }
            std::cout << "\n";

            std::cout << "d_hits:\n";
            for (int i = 0; i < numHits; i++) {
                std::cout << d_hits[i] << " ";
            }
            std::cout << "\n";

            std::cout << "\n";
            std::cout << "\n";
        }

        {
            std::cout << "with custom retrieve kernel, thread 0 prints and block synchronizes before retrieve func\n";
            thrust::device_vector<Location> d_hits(numHits);
            thrust::device_vector<CustomRetrieveKeyType> d_customKeysOfHits(numHits);

            thrust::device_vector<int> d_numKeysPerReadPrefixSum({0, numKeys});
            thrust::device_vector<int> d_numHitsPerReadPrefixSum({0, numHits});

            constexpr int bugtype = 2;
            callRetrieveHitsKernel<bugtype>(gpuHashTable, customKeys,
                numKeys, probe_eq,
                probe_doublehasher,
                d_numKeysPerReadPrefixSum.begin(),
                d_numHitsPerReadPrefixSum.begin(),
                numReads, 
                cuda::make_transform_output_iterator(
                    d_hits.begin(),
                    [] __device__(auto cucopair) -> Location { return cucopair.second; }
                ),
                d_customKeysOfHits.begin(),
                (cudaStream_t)0
            );

            cudaDeviceSynchronize();

            std::cout << "d_customKeysOfHits:\n";
            for (int i = 0; i < numHits; i++) {
                CustomRetrieveKeyType tup = d_customKeysOfHits[i];
                std::cout << "(" << cuda::std::get<0>(tup) << ","
                        << cuda::std::get<1>(tup) << ") ";
            }
            std::cout << "\n";

            std::cout << "d_hits:\n";
            for (int i = 0; i < numHits; i++) {
                std::cout << d_hits[i] << " ";
            }
            std::cout << "\n";

            std::cout << "\n";
            std::cout << "\n";
        }
    }

    std::cout << "\n";
    std::cout << "\n";



}

[sleeepyjack]

@fkallen my spider senses are tingling. There is a warp-collective right at the beginning of the cooperative retrieve(). With the printf branch only active for lane 0, this means it is called from a divergent context, and ITS is what stops it from silently working. Can you add a __syncwarp() right before the call to map.retrieve() and see if it magically resolves the issue?

[fkallen]

Adding __syncwarp() does not have a visible effect.