chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) 2021, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+#include <nv_util.h>
+
+#define TASK_PER_WARP_TILE_MACRO 1
+
+namespace gpu_cache {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+// GPU Cache API
+template <typename key_type>
+class gpu_cache_api {
+ public:
+  virtual ~gpu_cache_api() noexcept(false) {}
+  // Query API, i.e. A single read from the cache
+  virtual void Query(const key_type* d_keys, const size_t len, float* d_values,
+                     uint64_t* d_missing_index, key_type* d_missing_keys, size_t* d_missing_len,
+                     cudaStream_t stream,
+                     const size_t task_per_warp_tile = TASK_PER_WARP_TILE_MACRO) = 0;
+
+  // Replace API, i.e. Follow the Query API to update the content of the cache to Most Recent
+  virtual void Replace(const key_type* d_keys, const size_t len, const float* d_values,
+                       cudaStream_t stream,
+                       const size_t task_per_warp_tile = TASK_PER_WARP_TILE_MACRO) = 0;
+
+  // Update API, i.e. update the embeddings which exist in the cache
+  virtual void Update(const key_type* d_keys, const size_t len, const float* d_values,
+                      cudaStream_t stream,
+                      const size_t task_per_warp_tile = TASK_PER_WARP_TILE_MACRO) = 0;
+
+  // Dump API, i.e. dump some slabsets' keys from the cache
+  virtual void Dump(key_type* d_keys, size_t* d_dump_counter, const size_t start_set_index,
+                    const size_t end_set_index, cudaStream_t stream) = 0;
+
+  // Record all the lookup stream of a specific cache for Update/Replace sync
+  virtual void Record(cudaStream_t stream) = 0;
+};
+
+}  // namespace gpu_cache
@@ -0,0 +1,131 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cstdint>
+
+// MurmurHash3_32 implementation from
+// https://github.com/aappleby/smhasher/blob/master/src/MurmurHash3.cpp
+//-----------------------------------------------------------------------------
+// MurmurHash3 was written by Austin Appleby, and is placed in the public
+// domain. The author hereby disclaims copyright to this source code.
+// Note - The x86 and x64 versions do _not_ produce the same results, as the
+// algorithms are optimized for their respective platforms. You can still
+// compile and run any of them on any platform, but your performance with the
+// non-native version will be less than optimal.
+template <typename Key, uint32_t m_seed = 0>
+struct MurmurHash3_32 {
+  using argument_type = Key;
+  using result_type = uint32_t;
+
+  /*__forceinline__
+  __host__ __device__
+  MurmurHash3_32() : m_seed( 0 ) {}*/
+
+  __forceinline__ __host__ __device__ static uint32_t rotl32(uint32_t x, int8_t r) {
+    return (x << r) | (x >> (32 - r));
+  }
+
+  __forceinline__ __host__ __device__ static uint32_t fmix32(uint32_t h) {
+    h ^= h >> 16;
+    h *= 0x85ebca6b;
+    h ^= h >> 13;
+    h *= 0xc2b2ae35;
+    h ^= h >> 16;
+    return h;
+  }
+
+  /* --------------------------------------------------------------------------*/
+  /**
+   * @Synopsis  Combines two hash values into a new single hash value. Called
+   * repeatedly to create a hash value from several variables.
+   * Taken from the Boost hash_combine function
+   * https://www.boost.org/doc/libs/1_35_0/doc/html/boost/hash_combine_id241013.html
+   *
+   * @Param lhs The first hash value to combine
+   * @Param rhs The second hash value to combine
+   *
+   * @Returns A hash value that intelligently combines the lhs and rhs hash values
+   */
+  /* ----------------------------------------------------------------------------*/
+  __host__ __device__ static result_type hash_combine(result_type lhs, result_type rhs) {
+    result_type combined{lhs};
+
+    combined ^= rhs + 0x9e3779b9 + (combined << 6) + (combined >> 2);
+
+    return combined;
+  }
+
+  __forceinline__ __host__ __device__ static result_type hash(const Key& key) {
+    constexpr int len = sizeof(argument_type);
+    const uint8_t* const data = (const uint8_t*)&key;
+    constexpr int nblocks = len / 4;
+    uint32_t h1 = m_seed;
+    constexpr uint32_t c1 = 0xcc9e2d51;
+    constexpr uint32_t c2 = 0x1b873593;
+    //----------
+    // body
+    const uint32_t* const blocks = (const uint32_t*)(data + nblocks * 4);
+    for (int i = -nblocks; i; i++) {
+      uint32_t k1 = blocks[i];  // getblock32(blocks,i);
+      k1 *= c1;
+      k1 = rotl32(k1, 15);
+      k1 *= c2;
+      h1 ^= k1;
+      h1 = rotl32(h1, 13);
+      h1 = h1 * 5 + 0xe6546b64;
+    }
+    //----------
+    // tail
+    const uint8_t* tail = (const uint8_t*)(data + nblocks * 4);
+    uint32_t k1 = 0;
+    switch (len & 3) {
+      case 3:
+        k1 ^= tail[2] << 16;
+      case 2:
+        k1 ^= tail[1] << 8;
+      case 1:
+        k1 ^= tail[0];
+        k1 *= c1;
+        k1 = rotl32(k1, 15);
+        k1 *= c2;
+        h1 ^= k1;
+    };
+    //----------
+    // finalization
+    h1 ^= len;
+    h1 = fmix32(h1);
+    return h1;
+  }
+
+  __host__ __device__ __forceinline__ result_type operator()(const Key& key) const {
+    return this->hash(key);
+  }
+};
+
+template <typename key_type, typename index_type, index_type result>
+struct Fix_Hash {
+  using result_type = index_type;
+
+  __forceinline__ __host__ __device__ static index_type hash(const key_type& key) { return result; }
+};
+
+template <typename key_type, typename result_type>
+struct Mod_Hash {
+  __forceinline__ __host__ __device__ static result_type hash(const key_type& key) {
+    return (result_type)key;
+  }
+};
@@ -0,0 +1,122 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <nv_util.h>
+
+#include <cstdio>
+#include <hash_functions.cuh>
+#include <limits>
+
+#include "gpu_cache_api.hpp"
+#ifdef LIBCUDACXX_VERSION
+#include <cuda/atomic>
+#include <cuda/semaphore>
+#endif
+
+#define SET_ASSOCIATIVITY 2
+#define SLAB_SIZE 32
+#define TASK_PER_WARP_TILE_MACRO 1
+
+namespace gpu_cache {
+
+// slab for static slab list
+template <typename key_type, int warp_size>
+struct static_slab {
+  key_type slab_[warp_size];
+};
+
+// Static slablist(slabset) for GPU Cache
+template <int set_associativity, typename key_type, int warp_size>
+struct slab_set {
+  static_slab<key_type, warp_size> set_[set_associativity];
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+// GPU Cache
+template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
+          int warp_size, typename set_hasher = MurmurHash3_32<key_type>,
+          typename slab_hasher = Mod_Hash<key_type, size_t>>
+class gpu_cache : public gpu_cache_api<key_type> {
+ public:
+  // Ctor
+  gpu_cache(const size_t capacity_in_set, const size_t embedding_vec_size);
+
+  // Dtor
+  ~gpu_cache();
+
+  // Query API, i.e. A single read from the cache
+  void Query(const key_type* d_keys, const size_t len, float* d_values, uint64_t* d_missing_index,
+             key_type* d_missing_keys, size_t* d_missing_len, cudaStream_t stream,
+             const size_t task_per_warp_tile = TASK_PER_WARP_TILE_MACRO) override;
+
+  // Replace API, i.e. Follow the Query API to update the content of the cache to Most Recent
+  void Replace(const key_type* d_keys, const size_t len, const float* d_values, cudaStream_t stream,
+               const size_t task_per_warp_tile = TASK_PER_WARP_TILE_MACRO) override;
+
+  // Update API, i.e. update the embeddings which exist in the cache
+  void Update(const key_type* d_keys, const size_t len, const float* d_values, cudaStream_t stream,
+              const size_t task_per_warp_tile = TASK_PER_WARP_TILE_MACRO) override;
+
+  // Dump API, i.e. dump some slabsets' keys from the cache
+  void Dump(key_type* d_keys, size_t* d_dump_counter, const size_t start_set_index,
+            const size_t end_set_index, cudaStream_t stream) override;
+
+  void Record(cudaStream_t stream) override {}
+
+ public:
+  using slabset = slab_set<set_associativity, key_type, warp_size>;
+#ifdef LIBCUDACXX_VERSION
+  using atomic_ref_counter_type = cuda::atomic<ref_counter_type, cuda::thread_scope_device>;
+  using mutex = cuda::binary_semaphore<cuda::thread_scope_device>;
+#endif
+
+ private:
+  static const size_t BLOCK_SIZE_ = 64;
+
+  // Cache data
+  slabset* keys_;
+  float* vals_;
+  ref_counter_type* slot_counter_;
+
+  // Global counter
+#ifdef LIBCUDACXX_VERSION
+  atomic_ref_counter_type* global_counter_;
+#else
+  ref_counter_type* global_counter_;
+#endif
+  // CUDA device
+  int dev_;
+
+  // Cache capacity
+  size_t capacity_in_set_;
+  size_t num_slot_;
+
+  // Embedding vector size
+  size_t embedding_vec_size_;
+
+#ifdef LIBCUDACXX_VERSION
+  // Array of mutex to protect (sub-)warp-level data structure, each mutex protect 1 slab set
+  mutex* set_mutex_;
+#else
+  // Array of flag to protect (sub-)warp-level data structure, each flag act as a mutex and protect
+  // 1 slab set 1 for unlock, 0 for lock
+  int* set_mutex_;
+#endif
+};
+
+}  // namespace gpu_cache
@@ -0,0 +1,90 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime_api.h>
+
+#include <stdexcept>
+#include <string>
+
+#define CUDA_CHECK(val) \
+  { nv::cuda_check_((val), __FILE__, __LINE__); }
+
+namespace nv {
+
+template <typename T>
+struct is_fp8 : std::false_type {};
+
+template <>
+struct is_fp8<__nv_fp8_e4m3> : std::true_type {};
+
+template <>
+struct is_fp8<__nv_fp8_e5m2> : std::true_type {};
+
+class CudaException : public std::runtime_error {
+ public:
+  CudaException(const std::string& what) : runtime_error(what) {}
+};
+
+inline void cuda_check_(cudaError_t val, const char* file, int line) {
+  if (val != cudaSuccess) {
+    throw CudaException(std::string(file) + ":" + std::to_string(line) + ": CUDA error " +
+                        std::to_string(val) + ": " + cudaGetErrorString(val));
+  }
+}
+
+class CudaDeviceRestorer {
+ public:
+  CudaDeviceRestorer() { CUDA_CHECK(cudaGetDevice(&dev_)); }
+  ~CudaDeviceRestorer() { CUDA_CHECK(cudaSetDevice(dev_)); }
+  void check_device(int device) const {
+    if (device != dev_) {
+      throw std::runtime_error(
+          std::string(__FILE__) + ":" + std::to_string(__LINE__) +
+          ": Runtime Error: The device id in the context is not consistent with configuration");
+    }
+  }
+
+ private:
+  int dev_;
+};
+
+inline int get_dev(const void* ptr) {
+  cudaPointerAttributes attr;
+  CUDA_CHECK(cudaPointerGetAttributes(&attr, ptr));
+  int dev = -1;
+
+#if CUDART_VERSION >= 10000
+  if (attr.type == cudaMemoryTypeDevice)
+#else
+  if (attr.memoryType == cudaMemoryTypeDevice)
+#endif
+  {
+    dev = attr.device;
+  }
+  return dev;
+}
+
+inline void switch_to_dev(const void* ptr) {
+  int dev = get_dev(ptr);
+  if (dev >= 0) {
+    CUDA_CHECK(cudaSetDevice(dev));
+  }
+}
+
+}  // namespace nv