Add SSD backend integration and metadata API to DramKVEmbeddingCache

Lizhe Ji · facebook-github-bot · commit fa1da145c1cc · 2026-06-18T16:13:50.000-07:00
Summary: X-link: facebookresearch/FBGEMM#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
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
@@ -23,8 +23,10 @@
 #include <thrift/lib/cpp2/protocol/Serializer.h>
 #include <torch/script.h>
 #include <cmath>
+#include <cstring>
 #include <random>
 #include <string_view>
+#include <unordered_set>
 #include "common/time/Time.h"
 
 #include "../ssd_split_embeddings_cache/initializer.h"
@@ -369,6 +371,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<int64_t>(FixedBlockPool::get_metaheader_dim<weight_type>());
+    auto metadata_tensor = at::zeros(
+        {numel, metadata_dim},
+        at::TensorOptions().dtype(
+            c10::CppTypeToScalarType<weight_type>::value));
+    auto shardid_to_indexes = shard_input(indices, count);
+    std::vector<folly::Future<folly::Unit>> 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<index_t>();
+                      auto* md_ptr =
+                          metadata_tensor
+                              .template mutable_data_ptr<weight_type>();
+                      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<const char*>(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 +546,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 +697,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 +800,12 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
                               weights_data_ptr + tensor_offset * stride,
                               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 existing blocks
                           if (feature_evict_config_.has_value() &&
                               feature_evict_config_.value()->trigger_mode_ !=
@@ -765,6 +836,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_ !=
@@ -1781,6 +1857,10 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
 
   void compact() override {}
 
+  void set_ssd_backend() {
+    enable_ssd_backend_ = true;
+  }
+
   void trigger_feature_evict() {
     if (feature_evict_) {
       feature_evict_->trigger_evict();
@@ -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<weight_type>* 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<bool> enable_ssd_backend_{false};
 }; // class DramKVEmbeddingCache
 
 } // namespace kv_mem
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
@@ -0,0 +1,214 @@
+/*
+ * 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 <fmt/format.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <vector>
+
+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<DramKVEmbeddingCache<float>>(
+        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();
+  }
+
+  void insertEmbedding(int64_t id, float value = 1.0f) {
+    auto indices = at::tensor({id}, at::kLong);
+    std::vector<float> emb(EMBEDDING_DIM, value);
+    auto weights = at::from_blob(
+        emb.data(), {1, EMBEDDING_DIM}, at::TensorOptions().dtype(at::kFloat));
+    auto count = at::tensor({1}, at::kLong);
+    folly::coro::blockingWait(
+        dram_cache_->set_kv_db_async(indices, weights.clone(), count));
+  }
+
+  void insertEmbeddings(const std::vector<int64_t>& ids, float value = 1.0f) {
+    auto num = static_cast<int64_t>(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(
+        dram_cache_->set_kv_db_async(indices, weights, count));
+  }
+
+  std::shared_ptr<DramKVEmbeddingCache<float>> 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;
+  insertEmbedding(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<int64_t>(FixedBlockPool::get_metaheader_dim<float>());
+  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<float>(), 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<int64_t> keys = {1, 2, 3, 10, 100, 1000};
+  insertEmbeddings(keys, 1.0f);
+
+  auto indices = at::tensor(keys, at::kLong);
+  auto count = at::tensor({static_cast<int64_t>(keys.size())}, at::kLong);
+  auto metadata = dram_cache_->get_kv_metadata_rows(indices, count);
+
+  const int64_t expected_dim =
+      static_cast<int64_t>(FixedBlockPool::get_metaheader_dim<float>());
+  EXPECT_EQ(
+      metadata.sizes(),
+      at::IntArrayRef({static_cast<int64_t>(keys.size()), expected_dim}));
+
+  auto* md_ptr = metadata.data_ptr<float>();
+  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<int64_t>(FixedBlockPool::get_metaheader_dim<float>());
+  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;
+  insertEmbedding(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<float>(), 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
+  insertEmbedding(test_id, 3.0f);
+  auto metadata2 = dram_cache_->get_kv_metadata_rows(indices, count);
+  MetaHeader h2{};
+  std::memcpy(&h2, metadata2.data_ptr<float>(), 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<DramKVEmbeddingCache<at::Half>>(
+      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<int64_t>(FixedBlockPool::get_metaheader_dim<at::Half>());
+  // 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<at::Half>(), sizeof(int64_t));
+  EXPECT_EQ(decoded_key, 5);
+}
+
+} // namespace kv_mem
