Searched defs:EmbeddingDeviceCache (Results 1 – 1 of 1) sorted by relevance
56 struct EmbeddingDeviceCache { struct57 EmbeddingDeviceCache(size_t batch_elements, size_t cache_vocab_size) in EmbeddingDeviceCache() function69 std::unique_ptr<int[]> device_to_host_index;70 std::unique_ptr<int[]> device_to_host_ids;71 std::unique_ptr<int[]> host_to_device_index;72 std::unique_ptr<int[]> host_to_device_ids;73 int *hash_swap_index_addr_;74 float *hash_swap_value_addr_;75 std::shared_ptr<EmbeddingHashMap> device_hash_map_;76 std::shared_ptr<PsCacheBasic> cache_;