From 4d3d5d8b309baf2a2867127dad6b946140dc448c Mon Sep 17 00:00:00 2001 From: Lizhe Ji Date: Mon, 22 Jun 2026 16:53:49 -0700 Subject: [PATCH] Add SSD backend integration and metadata API to DramKVEmbeddingCache (#5929) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2847 Integrates `DramKVEmbeddingCache` with an SSD backend by exposing metadata retrieval APIs and internal state accessors. This enables the SSD tier to track dirty memory blocks for flushing, manage cross-tier feature eviction, and allows the enrichment process to skip redundant external data source fetches for IDs already present in SSD. Differential Revision: D108959007 --- .../dram_kv_embedding_cache.h | 113 ++++- .../fixed_block_pool.h | 7 - .../dram_kv_embedding_cache_test.cpp | 480 ++++++++++++++++++ .../fixed_block_pool_test.cpp | 5 +- 4 files changed, 593 insertions(+), 12 deletions(-) create mode 100644 fbgemm_gpu/test/dram_kv_embedding_cache/dram_kv_embedding_cache_test.cpp diff --git a/fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache.h b/fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache.h index 5af84a6a15..66ffc18cc2 100644 --- a/fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache.h +++ b/fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache.h @@ -23,8 +23,10 @@ #include #include #include +#include #include #include +#include #include "common/time/Time.h" #include "../ssd_split_embeddings_cache/initializer.h" @@ -103,6 +105,12 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB { /// @param enable_async_update whether to enable async update for the cache /// @param table_dims the table dimension for each table /// @param hash_size_cumsum the hash size cumulative sum for each table + /// @param enable_ssd_backend whether an SSD backend is wired up, enabling + /// dirty-bit tracking so dirty DRAM blocks can be written back to SSD. The + /// dirty bit denotes a change to the embedding value itself, as opposed to + /// its metadata; in the embedding-cache use case this only happens when a new + /// embedding is inserted, since the cache never updates an existing + /// embedding's value. Defaults to false. /// @return None explicit DramKVEmbeddingCache( int64_t max_D, @@ -126,7 +134,8 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB { std::vector table_offsets = {}, std::vector table_sizes = {}, std::optional> - enrichment_config = std::nullopt) + enrichment_config = std::nullopt, + bool enable_ssd_backend = false) : kv_db::EmbeddingKVDB( num_shards, max_D, @@ -150,12 +159,14 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB { num_shards_, block_size_, block_alignment_, - /*blocks_per_chunk=*/8192)), + /*blocks_per_chunk=*/8192, + /*enable_dirty_tracking=*/enable_ssd_backend)), elem_size_(row_storage_bitwidth / 8), backend_return_whole_row_(backend_return_whole_row), feature_evict_config_(std::move(feature_evict_config)), is_training_(is_training), - enable_raw_embedding_streaming_(enable_raw_embedding_streaming) { + enable_raw_embedding_streaming_(enable_raw_embedding_streaming), + enable_ssd_backend_(enable_ssd_backend) { executor_ = std::make_unique(std::max( num_threads, facebook::Proc::getCpuInfo().numCpuCores)); // Dedicated executor for enrichment (low priority, won't affect @@ -369,6 +380,62 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB { return metadata_tensor; } + at::Tensor get_kv_metadata_rows( + const at::Tensor& indices, + const at::Tensor& count) { + auto numel = indices.size(0); + const int64_t metadata_dim = + static_cast(FixedBlockPool::get_metaheader_dim()); + auto metadata_tensor = at::zeros( + {numel, metadata_dim}, + at::TensorOptions().dtype( + c10::CppTypeToScalarType::value)); + auto shardid_to_indexes = shard_input(indices, count); + std::vector> futures; + futures.reserve(shardid_to_indexes.size()); + const size_t metadata_bytes = metadata_dim * sizeof(weight_type); + for (const auto& [shard_id, indexes] : shardid_to_indexes) { + futures.emplace_back( + folly::via(executor_.get()) + .thenValue([this, + shard_id, + indexes, + &indices, + &metadata_tensor, + metadata_bytes](folly::Unit) { + FBGEMM_DISPATCH_INTEGRAL_TYPES( + indices.scalar_type(), + "dram_kv_metadata_rows", + [this, + shard_id, + indexes, + &indices, + &metadata_tensor, + metadata_bytes] { + using index_t = scalar_t; + CHECK(indices.is_contiguous()); + auto* idx_ptr = indices.const_data_ptr(); + auto* md_ptr = + metadata_tensor + .template mutable_data_ptr(); + const int64_t md_stride = metadata_tensor.size(1); + auto rlmap = kv_store_.by(shard_id).rlock(); + for (const auto& id_index : indexes) { + auto id = int64_t(idx_ptr[id_index]); + auto it = rlmap->find(id); + CHECK(it != rlmap->end()); + std::memcpy( + md_ptr + id_index * md_stride, + reinterpret_cast(it->second), + metadata_bytes); + } + }); + })); + } + folly::collect(futures).wait(); + return metadata_tensor; + } + /// insert embeddings into kvstore. /// current underlying memory management is done through F14FastMap /// key value pair will be sharded into multiple shards to increase @@ -488,6 +555,10 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB { weights_data_ptr + id_index * stride, weights_data_ptr + (id_index + 1) * stride, data_ptr); + // TODO: skip FixedBlockPool set_dirty here. This + // DRAM_SSD embedding cache path only handles + // backfill, where data already exists in SSD, so + // marking dirty would trigger a redundant flush. local_write_cache_copy_total_duration += facebook::WallClockUtil::NowInUsecFast() - before_copy_ts; @@ -635,6 +706,9 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB { weights_data_ptr + id_index * stride, weights_data_ptr + (id_index + 1) * stride, data_ptr); + if (enable_ssd_backend_) { + pool->set_dirty(block, true); + } cursor++; // Check if we should pause and yield lock if (is_laser_write_interrupted()) { @@ -735,6 +809,7 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB { weights_data_ptr + tensor_offset * stride, weights_data_ptr + (tensor_offset + 1) * stride, data_ptr); + // update provided ts for existing blocks if (feature_evict_config_.has_value() && feature_evict_config_.value()->trigger_mode_ != @@ -765,6 +840,11 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB { weights_data_ptr + (tensor_offset + 1) * stride, data_ptr); + // TODO: skip FixedBlockPool set_dirty here. This + // DRAM_SSD embedding cache path only handles + // backfill, where data already exists in SSD, so + // marking dirty would trigger a redundant flush. + // update provided ts for new allocated blocks if (feature_evict_config_.has_value() && feature_evict_config_.value()->trigger_mode_ != @@ -1945,6 +2025,25 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB { backend_return_whole_row_ = backend_return_whole_row; } + /// Get the feature evict object for callback wiring. + /// Returns nullptr if feature eviction is disabled. + FeatureEvict* get_feature_evict() { + return feature_evict_.get(); + } + + /// Access the internal kv_store for flush iteration. + auto& get_kv_store() { + return kv_store_; + } + + int64_t get_num_shards() const { + return num_shards_; + } + + int64_t get_block_size() const { + return block_size_; + } + private: int64_t get_dim_from_index(int64_t weight_idx) const { if (sub_table_dims_.empty()) { @@ -2378,6 +2477,9 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB { weights_data_ptr + id_index * stride, weights_data_ptr + (id_index + 1) * stride, block); + if (enable_ssd_backend_) { + pool->set_dirty(block, true); + } if (new_block) { if (feature_evict_config_.has_value() && @@ -2501,6 +2603,11 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB { // OpenTab/Maple reader for ONEFLOW_OPENTAB_SID enrichment (type-erased) oneflow_enrichment::ReaderPtr open_tab_reader_; + + // Optional SSD backend for existence checks during enrichment. + // When set, enrichment will skip IDs that already exist in SSD, + // avoiding unnecessary calls to external data sources. + std::atomic enable_ssd_backend_{false}; }; // class DramKVEmbeddingCache } // namespace kv_mem diff --git a/fbgemm_gpu/src/dram_kv_embedding_cache/fixed_block_pool.h b/fbgemm_gpu/src/dram_kv_embedding_cache/fixed_block_pool.h index 89767d2ee6..fe038067d0 100644 --- a/fbgemm_gpu/src/dram_kv_embedding_cache/fixed_block_pool.h +++ b/fbgemm_gpu/src/dram_kv_embedding_cache/fixed_block_pool.h @@ -362,13 +362,6 @@ class FixedBlockPool : public std::pmr::memory_resource { set_used(result, true); set_count(result, 0); update_timestamp(result); - if (enable_dirty_tracking_) { - // Intentional: a freshly allocated block holds new data that has not yet - // been persisted to SSD, so it must be marked dirty to be picked up by - // the SSD flush protocol. Callers should not assume "allocate" returns - // clean memory. - set_dirty(result, true); - } return result; } diff --git a/fbgemm_gpu/test/dram_kv_embedding_cache/dram_kv_embedding_cache_test.cpp b/fbgemm_gpu/test/dram_kv_embedding_cache/dram_kv_embedding_cache_test.cpp new file mode 100644 index 0000000000..66c6011fd9 --- /dev/null +++ b/fbgemm_gpu/test/dram_kv_embedding_cache/dram_kv_embedding_cache_test.cpp @@ -0,0 +1,480 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include "deeplearning/fbgemm/fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache.h" + +#include +#include +#include +#include + +namespace kv_mem { + +struct MetaHeader { + int64_t key; + uint32_t timestamp; + uint32_t count : 31; + bool used : 1; +}; + +class DramKVEmbeddingCacheTest : public ::testing::Test { + protected: + static constexpr int EMBEDDING_DIM = 16; + static constexpr int NUM_SHARDS = 4; + + void SetUp() override { + FLAGS_logtostderr = true; + FLAGS_minloglevel = 0; + + auto hash_size_cumsum = at::tensor({0, 100000}, at::kLong); + + dram_cache_ = std::make_shared>( + EMBEDDING_DIM, + /*uniform_init_lower=*/-0.1, + /*uniform_init_upper=*/0.1, + /*feature_evict_config=*/std::nullopt, + NUM_SHARDS, + /*num_threads=*/4, + /*row_storage_bitwidth=*/32, + /*backend_return_whole_row=*/false, + /*enable_async_update=*/false, + /*table_dims=*/std::nullopt, + hash_size_cumsum, + /*is_training=*/false, + /*disable_random_init=*/true); + } + + void TearDown() override { + dram_cache_.reset(); + } + + // Thin wrappers named after the core cache APIs they exercise, so test bodies + // read like direct calls into the cache under test. + void set_kv_db_async( + DramKVEmbeddingCache& cache, + int64_t id, + float value = 1.0f) { + auto indices = at::tensor({id}, at::kLong); + auto weights = at::full( + {1, EMBEDDING_DIM}, value, at::TensorOptions().dtype(at::kFloat)); + auto count = at::tensor({1}, at::kLong); + folly::coro::blockingWait(cache.set_kv_db_async(indices, weights, count)); + } + + void set_kv_db_async( + DramKVEmbeddingCache& cache, + const std::vector& ids, + float value = 1.0f) { + auto num = static_cast(ids.size()); + auto indices = at::tensor(ids, at::kLong); + auto weights = at::full( + {num, EMBEDDING_DIM}, value, at::TensorOptions().dtype(at::kFloat)); + auto count = at::tensor({num}, at::kLong); + folly::coro::blockingWait(cache.set_kv_db_async(indices, weights, count)); + } + + // Build a standalone float cache, varying only enable_ssd_backend so tests + // can compare dirty-bit behavior with the flag on vs off. + std::shared_ptr> makeCache( + bool enable_ssd_backend) { + auto hash_size_cumsum = at::tensor({0, 100000}, at::kLong); + return std::make_shared>( + EMBEDDING_DIM, + /*uniform_init_lower=*/-0.1, + /*uniform_init_upper=*/0.1, + /*feature_evict_config=*/std::nullopt, + NUM_SHARDS, + /*num_threads=*/4, + /*row_storage_bitwidth=*/32, + /*backend_return_whole_row=*/false, + /*enable_async_update=*/false, + /*table_dims=*/std::nullopt, + hash_size_cumsum, + /*is_training=*/false, + /*disable_random_init=*/true, + /*enable_raw_embedding_streaming=*/false, + /*res_store_shards=*/0, + /*res_server_port=*/0, + /*table_names=*/std::vector{}, + /*table_offsets=*/std::vector{}, + /*table_sizes=*/std::vector{}, + /*enrichment_config=*/std::nullopt, + /*enable_ssd_backend=*/enable_ssd_backend); + } + + // Look up the block for `id` and report its dirty bit. The dirty state lives + // in the owning shard's pool, so query get_dirty() on that same pool. + bool isKeyDirty(DramKVEmbeddingCache& cache, int64_t id) { + auto& kv_store = cache.get_kv_store(); + for (int shard = 0; shard < cache.get_num_shards(); ++shard) { + auto rlmap = kv_store.by(shard).rlock(); + auto it = rlmap->find(id); + if (it != rlmap->end()) { + return kv_store.pool_by(shard)->get_dirty(it->second); + } + } + ADD_FAILURE() << "key " << id << " not found in any shard"; + return false; + } + + // Clear the dirty bit for `id` so a subsequent write must re-set it. + void clearDirty(DramKVEmbeddingCache& cache, int64_t id) { + auto& kv_store = cache.get_kv_store(); + for (int shard = 0; shard < cache.get_num_shards(); ++shard) { + auto rlmap = kv_store.by(shard).rlock(); + auto it = rlmap->find(id); + if (it != rlmap->end()) { + kv_store.pool_by(shard)->clear_dirty(it->second); + return; + } + } + ADD_FAILURE() << "key " << id << " not found to clear dirty"; + } + + // Write a single row through the metaheader storage path. The key is encoded + // into the first 8 bytes of the row, mirroring how checkpoint restore feeds + // weights-with-metaheader into the cache. + void set_kv_with_metaheader_to_storage( + DramKVEmbeddingCache& cache, + int64_t id) { + const int64_t metaheader_dim = + static_cast(FixedBlockPool::get_metaheader_dim()); + const int64_t total_width = metaheader_dim + EMBEDDING_DIM; + auto weights = + at::zeros({1, total_width}, at::TensorOptions().dtype(at::kFloat)); + FixedBlockPool::set_key(weights[0].data_ptr(), id); + cache.set_kv_with_metaheader_to_storage(weights); + } + + // inference_set_kv_db_async requires feature_evict_ to be a + // TimeThresholdBasedEvict (CHECK in the method), so build the cache with a + // BY_TIMESTAMP_THRESHOLD feature-evict config (MANUAL trigger needs no extra + // fields). + std::shared_ptr> makeCacheWithTimeThresholdEvict( + bool enable_ssd_backend) { + auto feature_evict_config = c10::make_intrusive( + /*trigger_mode=*/static_cast(EvictTriggerMode::MANUAL), + /*trigger_strategy=*/ + static_cast(EvictTriggerStrategy::BY_TIMESTAMP_THRESHOLD), + /*trigger_step_interval=*/std::nullopt, + /*mem_util_threshold_in_GB=*/std::nullopt, + /*ttls_in_mins=*/std::nullopt, + /*counter_thresholds=*/std::nullopt, + /*counter_decay_rates=*/std::nullopt, + /*feature_score_counter_decay_rates=*/std::nullopt, + /*training_id_eviction_trigger_count=*/std::nullopt, + /*training_id_keep_count=*/std::nullopt, + /*enable_eviction_for_feature_score_eviction_policy=*/std::nullopt, + /*l2_weight_thresholds=*/std::nullopt, + /*embedding_dims=*/std::nullopt); + auto hash_size_cumsum = at::tensor({0, 100000}, at::kLong); + return std::make_shared>( + EMBEDDING_DIM, + /*uniform_init_lower=*/-0.1, + /*uniform_init_upper=*/0.1, + feature_evict_config, + NUM_SHARDS, + /*num_threads=*/4, + /*row_storage_bitwidth=*/32, + /*backend_return_whole_row=*/false, + /*enable_async_update=*/false, + /*table_dims=*/std::nullopt, + hash_size_cumsum, + /*is_training=*/false, + /*disable_random_init=*/true, + /*enable_raw_embedding_streaming=*/false, + /*res_store_shards=*/0, + /*res_server_port=*/0, + /*table_names=*/std::vector{}, + /*table_offsets=*/std::vector{}, + /*table_sizes=*/std::vector{}, + /*enrichment_config=*/std::nullopt, + /*enable_ssd_backend=*/enable_ssd_backend); + } + + // Thin wrapper named after the core API, mirroring set_kv_db_async. + void inference_set_kv_db_async( + DramKVEmbeddingCache& cache, + int64_t id, + float value = 1.0f, + std::optional inplace_update_ts = std::nullopt) { + auto indices = at::tensor({id}, at::kLong); + auto weights = at::full( + {1, EMBEDDING_DIM}, value, at::TensorOptions().dtype(at::kFloat)); + auto count = at::tensor({1}, at::kLong); + folly::coro::blockingWait(cache.inference_set_kv_db_async( + indices, weights, count, inplace_update_ts)); + } + + // set_kv_zch_eviction_metadata_async is a no-op unless feature_evict_ is a + // FeatureScoreBasedEvict, so build the cache with a BY_FEATURE_SCORE config. + // hash_size_cumsum = {0, 100000} -> a single sub-table, so the per-sub-table + // vectors are size 1. + std::shared_ptr> makeCacheWithFeatureScoreEvict( + bool enable_ssd_backend) { + auto feature_evict_config = c10::make_intrusive( + /*trigger_mode=*/static_cast(EvictTriggerMode::MANUAL), + /*trigger_strategy=*/ + static_cast(EvictTriggerStrategy::BY_FEATURE_SCORE), + /*trigger_step_interval=*/std::nullopt, + /*mem_util_threshold_in_GB=*/std::nullopt, + /*ttls_in_mins=*/std::vector{1}, + /*counter_thresholds=*/std::nullopt, + /*counter_decay_rates=*/std::nullopt, + /*feature_score_counter_decay_rates=*/std::vector{0.5}, + /*training_id_eviction_trigger_count=*/std::vector{1000}, + /*training_id_keep_count=*/std::vector{100}, + /*enable_eviction_for_feature_score_eviction_policy=*/ + std::vector{1}, + /*l2_weight_thresholds=*/std::nullopt, + /*embedding_dims=*/std::nullopt); + auto hash_size_cumsum = at::tensor({0, 100000}, at::kLong); + return std::make_shared>( + EMBEDDING_DIM, + /*uniform_init_lower=*/-0.1, + /*uniform_init_upper=*/0.1, + feature_evict_config, + NUM_SHARDS, + /*num_threads=*/4, + /*row_storage_bitwidth=*/32, + /*backend_return_whole_row=*/false, + /*enable_async_update=*/false, + /*table_dims=*/std::nullopt, + hash_size_cumsum, + /*is_training=*/false, + /*disable_random_init=*/true, + /*enable_raw_embedding_streaming=*/false, + /*res_store_shards=*/0, + /*res_server_port=*/0, + /*table_names=*/std::vector{}, + /*table_offsets=*/std::vector{}, + /*table_sizes=*/std::vector{}, + /*enrichment_config=*/std::nullopt, + /*enable_ssd_backend=*/enable_ssd_backend); + } + + // Thin wrapper named after the core API, mirroring set_kv_db_async. + void set_kv_zch_eviction_metadata_async( + DramKVEmbeddingCache& cache, + int64_t id, + float engage_rate = 1.0f) { + auto indices = at::tensor({id}, at::kLong); + auto count = at::tensor({1}, at::kLong); + auto engage_rates = + at::full({1}, engage_rate, at::TensorOptions().dtype(at::kFloat)); + folly::coro::blockingWait( + cache.set_kv_zch_eviction_metadata_async(indices, count, engage_rates)); + } + + std::shared_ptr> dram_cache_; +}; + +// Test: get_kv_metadata_rows returns correct shape and key for single inserted +// id +TEST_F(DramKVEmbeddingCacheTest, SingleKeyMetadata) { + const int64_t test_id = 42; + set_kv_db_async(*dram_cache_, test_id, 2.5f); + + auto indices = at::tensor({test_id}, at::kLong); + auto count = at::tensor({1}, at::kLong); + auto metadata = dram_cache_->get_kv_metadata_rows(indices, count); + + const int64_t expected_dim = + static_cast(FixedBlockPool::get_metaheader_dim()); + EXPECT_EQ(metadata.dim(), 2); + EXPECT_EQ(metadata.size(0), 1); + EXPECT_EQ(metadata.size(1), expected_dim); + EXPECT_EQ(metadata.dtype(), at::kFloat); + static_assert(sizeof(MetaHeader) == 16, "MetaHeader must be 16 bytes"); + + MetaHeader header{}; + std::memcpy(&header, metadata.data_ptr(), sizeof(MetaHeader)); + + EXPECT_EQ(header.key, test_id); + EXPECT_TRUE(header.used); + EXPECT_GT(header.timestamp, 0u); + // count may be 0 initially or updated depending on implementation + EXPECT_GE(header.count, 0u); +} + +// Test: get_kv_metadata_rows returns correct metadata for multiple keys across +// shards +TEST_F(DramKVEmbeddingCacheTest, MultipleKeysMetadata) { + std::vector keys = {1, 2, 3, 10, 100, 1000}; + set_kv_db_async(*dram_cache_, keys, 1.0f); + + auto indices = at::tensor(keys, at::kLong); + auto count = at::tensor({static_cast(keys.size())}, at::kLong); + auto metadata = dram_cache_->get_kv_metadata_rows(indices, count); + + const int64_t expected_dim = + static_cast(FixedBlockPool::get_metaheader_dim()); + EXPECT_EQ( + metadata.sizes(), + at::IntArrayRef({static_cast(keys.size()), expected_dim})); + + auto* md_ptr = metadata.data_ptr(); + const int64_t stride = expected_dim; + for (size_t i = 0; i < keys.size(); ++i) { + MetaHeader header{}; + std::memcpy(&header, md_ptr + i * stride, sizeof(MetaHeader)); + EXPECT_EQ(header.key, keys[i]) << "Mismatch at index " << i; + EXPECT_TRUE(header.used) << "Used flag false for key " << keys[i]; + EXPECT_GT(header.timestamp, 0u) << "Timestamp not set for key " << keys[i]; + } +} + +// Test: get_kv_metadata_rows with empty input returns empty tensor with correct +// dim +TEST_F(DramKVEmbeddingCacheTest, EmptyInputReturnsEmpty) { + auto indices = at::empty({0}, at::kLong); + auto count = at::tensor({0}, at::kLong); + auto metadata = dram_cache_->get_kv_metadata_rows(indices, count); + + const int64_t expected_dim = + static_cast(FixedBlockPool::get_metaheader_dim()); + EXPECT_EQ(metadata.dim(), 2); + EXPECT_EQ(metadata.size(0), 0); + EXPECT_EQ(metadata.size(1), expected_dim); +} + +// Test: get_kv_metadata_rows reflects updated timestamp after re-insert +TEST_F(DramKVEmbeddingCacheTest, TimestampUpdatesOnReinsert) { + const int64_t test_id = 7; + set_kv_db_async(*dram_cache_, test_id, 1.0f); + + auto indices = at::tensor({test_id}, at::kLong); + auto count = at::tensor({1}, at::kLong); + auto metadata1 = dram_cache_->get_kv_metadata_rows(indices, count); + MetaHeader h1{}; + std::memcpy(&h1, metadata1.data_ptr(), sizeof(MetaHeader)); + + // Sleep to ensure timestamp advances (timestamp is in seconds) + std::this_thread::sleep_for(std::chrono::seconds(2)); + + // Re-insert same key to update timestamp + set_kv_db_async(*dram_cache_, test_id, 3.0f); + auto metadata2 = dram_cache_->get_kv_metadata_rows(indices, count); + MetaHeader h2{}; + std::memcpy(&h2, metadata2.data_ptr(), sizeof(MetaHeader)); + + EXPECT_EQ(h2.key, test_id); + EXPECT_TRUE(h2.used); + EXPECT_GE(h2.timestamp, h1.timestamp); +} + +// Test: get_kv_metadata_rows works with float16 weight type via separate cache +// instance +TEST_F(DramKVEmbeddingCacheTest, HalfPrecisionMetadataDim) { + auto hash_size_cumsum = at::tensor({0, 100000}, at::kLong); + auto dram_cache_half = std::make_shared>( + EMBEDDING_DIM, + -0.1, + 0.1, + std::nullopt, + NUM_SHARDS, + 4, + 16, + false, + false, + std::nullopt, + hash_size_cumsum, + false, + true); + + // Insert one key + auto indices = at::tensor({5}, at::kLong); + auto weights = + at::full({1, EMBEDDING_DIM}, 1.0, at::TensorOptions().dtype(at::kHalf)); + auto count = at::tensor({1}, at::kLong); + folly::coro::blockingWait( + dram_cache_half->set_kv_db_async(indices, weights, count)); + + auto metadata = dram_cache_half->get_kv_metadata_rows(indices, count); + const int64_t expected_dim = + static_cast(FixedBlockPool::get_metaheader_dim()); + // 16 bytes / 2 bytes per half = 8 + EXPECT_EQ(expected_dim, 8); + EXPECT_EQ(metadata.sizes(), at::IntArrayRef({1, expected_dim})); + EXPECT_EQ(metadata.dtype(), at::kHalf); + + // Decode first 8 bytes as int64 key from half tensor raw bytes + int64_t decoded_key = 0; + std::memcpy(&decoded_key, metadata.data_ptr(), sizeof(int64_t)); + EXPECT_EQ(decoded_key, 5); +} + +// Shared fixture for all enable_ssd_backend-parameterized dirty-bit tests, so a +// single on/off instantiation covers every TEST_P below. +class DirtyBitParamTest : public DramKVEmbeddingCacheTest, + public ::testing::WithParamInterface {}; + +// Test: the set_kv_db_async path always leaves the block clean because there is +// no explicit set_dirty call (and allocation no longer marks blocks dirty), +// regardless of enable_ssd_backend. +TEST_P(DirtyBitParamTest, SetKvDbAsyncLeavesBlockClean) { + const bool enable_ssd_backend = GetParam(); + auto cache = makeCache(enable_ssd_backend); + const int64_t test_id = 123; + set_kv_db_async(*cache, test_id); + + EXPECT_FALSE(isKeyDirty(*cache, test_id)); +} + +// Test: like set_kv_db_async, inference_set_kv_db_async leaves the block clean +// because it has no explicit set_dirty call (and allocation no longer marks +// blocks dirty), regardless of enable_ssd_backend. +TEST_P(DirtyBitParamTest, InferenceSetKvDbAsyncLeavesBlockClean) { + const bool enable_ssd_backend = GetParam(); + auto cache = makeCacheWithTimeThresholdEvict(enable_ssd_backend); + const int64_t test_id = 123; + inference_set_kv_db_async(*cache, test_id); + + EXPECT_FALSE(isKeyDirty(*cache, test_id)); +} + +// Test: the metaheader write path marks a block dirty iff enable_ssd_backend is +// set. The block is allocated and its dirty bit cleared first, so the assertion +// isolates the `if (enable_ssd_backend_) set_dirty(...)` branch from +// allocation-time state. +TEST_P(DirtyBitParamTest, MetaheaderWriteDirtiesIffBackendEnabled) { + const bool enable_ssd_backend = GetParam(); + auto cache = makeCache(enable_ssd_backend); + const int64_t test_id = 77; + + set_kv_with_metaheader_to_storage(*cache, test_id); + clearDirty(*cache, test_id); + ASSERT_FALSE(isKeyDirty(*cache, test_id)); + + set_kv_with_metaheader_to_storage(*cache, test_id); + + EXPECT_EQ(isKeyDirty(*cache, test_id), enable_ssd_backend); +} + +// Test: the feature-score metadata path (set_kv_zch_eviction_metadata_async) +// always leaves the block clean because there is no explicit set_dirty call +// (and allocation no longer marks blocks dirty), regardless of +// enable_ssd_backend. +TEST_P(DirtyBitParamTest, FeatureScoreMetadataDirtiesIffBackendEnabled) { + const bool enable_ssd_backend = GetParam(); + auto cache = makeCacheWithFeatureScoreEvict(enable_ssd_backend); + const int64_t test_id = 55; + set_kv_zch_eviction_metadata_async(*cache, test_id, /*engage_rate=*/0.5f); + + EXPECT_FALSE(isKeyDirty(*cache, test_id)); +} + +INSTANTIATE_TEST_SUITE_P( + EnableSsdBackend, + DirtyBitParamTest, + ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return info.param ? "On" : "Off"; + }); + +} // namespace kv_mem diff --git a/fbgemm_gpu/test/dram_kv_embedding_cache/fixed_block_pool_test.cpp b/fbgemm_gpu/test/dram_kv_embedding_cache/fixed_block_pool_test.cpp index 7dffd37c49..b3271c05dc 100644 --- a/fbgemm_gpu/test/dram_kv_embedding_cache/fixed_block_pool_test.cpp +++ b/fbgemm_gpu/test/dram_kv_embedding_cache/fixed_block_pool_test.cpp @@ -431,8 +431,6 @@ TEST(FixedBlockPool, DefaultDirtyTrackingGating) { FixedBlockPool default_pool(block_size, alignment, 1024); EXPECT_FALSE(default_pool.is_dirty_tracking_enabled()); - // After allocation: a freshly allocated block is always clean for disabled - // dirty tracking. auto* block = default_pool.allocate_t(); ASSERT_NE(block, nullptr); EXPECT_FALSE(default_pool.get_dirty(block)); @@ -460,6 +458,9 @@ TEST(FixedBlockPool, DirtyTrackingGating) { // yet been persisted to SSD). auto* block = tracked_pool.allocate_t(); ASSERT_NE(block, nullptr); + EXPECT_FALSE(tracked_pool.get_dirty(block)); + + tracked_pool.set_dirty(block, true); EXPECT_TRUE(tracked_pool.get_dirty(block)); // After operations: clearing (e.g. once flushed to SSD) marks it clean,