Add LFU strategy to dynamicemb

Runchu Zhao · Runchu Zhao · commit 1a354c926631 · 2025-05-25T11:01:11.000Z
diff --git a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py
@@ -605,6 +605,7 @@ def forward(
         self._update_score()
         return res
 
+    # Todo: add LFU here, LFU really need to set score?
     def set_score(
         self,
         named_score: Dict[str, int],
@@ -646,6 +647,8 @@ def _create_score(self):
                 self._scores[table_name] = 1
             elif option.score_strategy == DynamicEmbScoreStrategy.CUSTOMIZED:
                 option.evict_strategy = DynamicEmbEvictStrategy.CUSTOMIZED
+            elif option.score_strategy == DynamicEmbScoreStrategy.LFU:
+                option.evict_strategy = DynamicEmbEvictStrategy.LFU #ADD LFU here
 
     def _update_score(self):
         for table_name, option in zip(self._table_names, self._dynamicemb_options):
@@ -669,6 +672,9 @@ def _update_score(self):
                     self._scores[table_name] = 0
                 else:
                     self._scores[table_name] = new_score
+            elif option.score_strategy == DynamicEmbScoreStrategy.LFU:
+                new_score = old_score + 1 # Todo : confirm score update logic
+                self._scores[table_name] = new_score + old_score
 
     def incremental_dump(
         self,
diff --git a/corelib/dynamicemb/dynamicemb/dynamicemb_config.py b/corelib/dynamicemb/dynamicemb/dynamicemb_config.py
@@ -199,7 +199,7 @@ class DynamicEmbScoreStrategy(enum.IntEnum):
     TIMESTAMP = 0
     STEP = 1
     CUSTOMIZED = 2
-
+    LFU = 3 #ADD LFU here
 
 # Configs used as keys to group HKV variables(considering kernel behaviors, result type).
 @dataclass
diff --git a/corelib/dynamicemb/src/dynamic_emb_op.cu b/corelib/dynamicemb/src/dynamic_emb_op.cu
@@ -138,11 +138,11 @@ void insert_and_evict(
     bool unique_key = true,
     bool ignore_evict_strategy = false) {
 
-  if (not score and table->evict_strategy() == EvictStrategy::kCustomized) {
-    throw std::invalid_argument("Must specify the score when evict strategy is customized.");
+  if (not score and (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu)) {
+    throw std::invalid_argument("Must specify the score when evict strategy is customized or LFU.");
   }
   auto stream = at::cuda::getCurrentCUDAStream().stream();
-  if (table->evict_strategy() == EvictStrategy::kCustomized) {
+  if (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu) {
     auto&& option = at::TensorOptions().dtype(at::kUInt64).device(keys.device());
     // broadcast scores
     at::Tensor bc_scores = at::empty({static_cast<int64_t>(n)}, option);
@@ -188,8 +188,8 @@ void find_or_insert(std::shared_ptr<dyn_emb::DynamicVariableBase> table,
                   bool ignore_evict_strategy = false
                   )
 {
-  if (not score and table->evict_strategy() == EvictStrategy::kCustomized) {
-    throw std::invalid_argument("Must specify the score when evict strategy is customized.");
+  if (not score and (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu)) {
+    throw std::invalid_argument("Must specify the score when evict strategy is customized or LFU.");
   }
   if (n == 0) return;
   auto stream = at::cuda::getCurrentCUDAStream().stream();
@@ -203,7 +203,7 @@ void find_or_insert(std::shared_ptr<dyn_emb::DynamicVariableBase> table,
 
   auto found_tensor_data_ptr = found_tensor.data_ptr<bool>();
 
-  if (table->evict_strategy() == EvictStrategy::kCustomized) {
+  if (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu) {
     auto&& option = at::TensorOptions().dtype(at::kUInt64).device(keys.device());
     // broadcast scores
     at::Tensor bc_scores = at::empty({static_cast<int64_t>(n)}, option);
@@ -226,15 +226,15 @@ void find_or_insert_pointers(
   const std::optional<uint64_t> score = std::nullopt,
   bool unique_key = true,
   bool ignore_evict_strategy = false) {
-  if (not score and table->evict_strategy() == EvictStrategy::kCustomized) {
-    throw std::invalid_argument("Must specify the score when evict strategy is customized.");
+  if (not score and (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu)) {
+    throw std::invalid_argument("Must specify the score when evict strategy is customized or LFU.");
   }
   if (n == 0) return;
   auto stream = at::cuda::getCurrentCUDAStream().stream();
   auto values_data_ptr = reinterpret_cast<void**>(values.data_ptr<int64_t>());
   auto found_tensor_data_ptr = founds.data_ptr<bool>();
 
-  if (table->evict_strategy() == EvictStrategy::kCustomized) {
+  if (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu) {
     auto&& option = at::TensorOptions().dtype(at::kUInt64).device(keys.device());
     // broadcast scores
     at::Tensor bc_scores = at::empty({static_cast<int64_t>(n)}, option);
diff --git a/corelib/dynamicemb/src/dynamic_variable_base.cu b/corelib/dynamicemb/src/dynamic_variable_base.cu
@@ -29,6 +29,7 @@ std::shared_ptr<DynamicVariableBase> VariableFactory::create(
     const SafeCheckMode safe_check_mode = SafeCheckMode::IGNORE) {
   // TODO:check input datetype , now only support key type is int64_t uint64_t ,
   // value type float
+  printf("create dynamic variable base\n");
   std::shared_ptr<DynamicVariableBase> table;
   DISPATCH_INTEGER_DATATYPE_FUNCTION(key_type, key_t, [&] {
     DISPATCH_FLOAT_DATATYPE_FUNCTION(value_type, value_t, [&] {
diff --git a/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__int64_t___half_EvictStrategy__kLfu.cu b/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__int64_t___half_EvictStrategy__kLfu.cu
@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<int64_t, __half, EvictStrategy::kLfu>;
+} 
diff --git a/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__int64_t___nv_bfloat16_EvictStrategy__kLfu.cu b/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__int64_t___nv_bfloat16_EvictStrategy__kLfu.cu
@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<int64_t, __nv_bfloat16, EvictStrategy::kLfu>;
+} 
diff --git a/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__int64_t_float_EvictStrategy__kLfu.cu b/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__int64_t_float_EvictStrategy__kLfu.cu
@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<int64_t, float, EvictStrategy::kLfu>;
+} 
diff --git a/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__uint64_t___half_EvictStrategy__kLfu.cu b/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__uint64_t___half_EvictStrategy__kLfu.cu
@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<uint64_t, __half, EvictStrategy::kLfu>;
+} 
diff --git a/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__uint64_t___nv_bfloat16_EvictStrategy__kLfu.cu b/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__uint64_t___nv_bfloat16_EvictStrategy__kLfu.cu
@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<uint64_t, __nv_bfloat16, EvictStrategy::kLfu>;
+} 
diff --git a/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__uint64_t_float_EvictStrategy__kLfu.cu b/corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__uint64_t_float_EvictStrategy__kLfu.cu
@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<uint64_t, float, EvictStrategy::kLfu>;
+} 
diff --git a/corelib/dynamicemb/src/utils.h b/corelib/dynamicemb/src/utils.h
@@ -104,6 +104,7 @@ enum class EvictStrategy : uint32_t {
   switch (EVICT_TYPE) {                                                 \
     CASE_ENUM_USING_HINT(EvictStrategy::kLru, HINT, __VA_ARGS__)        \
     CASE_ENUM_USING_HINT(EvictStrategy::kCustomized, HINT, __VA_ARGS__) \
+    CASE_ENUM_USING_HINT(EvictStrategy::kLfu, HINT, __VA_ARGS__) \
     default:                                                            \
       exit(EXIT_FAILURE);                                               \
   }
diff --git a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_dynamicemb_extensions.py b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_dynamicemb_extensions.py
@@ -65,28 +65,42 @@ def __init__(
         self.device_ = device
 
         min_step = 0
-        self.min_score_: int = (
-            device_timestamp() if evict_strategy == EvictStrategy.KLru else min_step
-        )
+        if evict_strategy == EvictStrategy.KLru:
+            self.min_score_: int = device_timestamp()
+        elif evict_strategy == EvictStrategy.KLfu:
+            # LFU 模式：使用单调递增的分数
+            self.min_score_: int = min_step
+        else:
+            self.min_score_: int = min_step
         self.step_: int = min_step + 1
 
     def min_score(self):
         return self.min_score_
 
+    #   add LFU here same to customized
     def next_score(self) -> int:
-        next_score = self.step_
-        return (
-            device_timestamp()
-            if self.evict_strategy_ == EvictStrategy.KLru
-            else next_score
-        )
-
+        if self.evict_strategy_ == EvictStrategy.KLru:
+            return device_timestamp()
+        else:
+            # LFU 和其他模式：返回单调递增的分数
+            next_score = self.step_
+            self.step_ += 1
+            return next_score
+
+    # add LFU here
     def score(self):
         if self.evict_strategy_ == EvictStrategy.KLru:
             return None
-        score = self.step_
-        self.step_ += 1
-        return score
+        elif self.evict_strategy_ == EvictStrategy.KLfu:
+            # score = self.step_
+            # self.step_ += 1
+            # return score
+            # LFU 应给merlin hashtable频率增量1 
+            return 1 
+        else:
+            score = self.step_
+            self.step_ += 1
+            return score
 
 
 def random_indices(batch, min_index, max_index):
@@ -125,9 +139,12 @@ def score_type():
 def counter_dtype():
     return torch.uint64
 
-
+# add LFU here
+# @pytest.mark.parametrize(
+#     "evict_strategy", [EvictStrategy.KLru, EvictStrategy.KCustomized, EvictStrategy.KLfu]
+# )
 @pytest.mark.parametrize(
-    "evict_strategy", [EvictStrategy.KLru, EvictStrategy.KCustomized]
+    "evict_strategy", [EvictStrategy.KLfu]
 )
 @pytest.mark.parametrize(
     "bucket_capacity, batch, capacity, num_iteration, dump_interval",
@@ -168,6 +185,7 @@ def test_dynamicemb_extensions(
     assert ext_option.dim * ext_option.max_capacity <= ext_option.local_hbm_for_values
 
     table = DynamicEmbTable(*astuple(ext_option))
+    # import pdb; pdb.set_trace() 
     device = torch.device(f"cuda:{ext_option.device_id}")
     score_adaptor = ScoreAdaptor(evict_strategy, score_type, device)
     init_score: int = score_adaptor.min_score()
@@ -245,6 +263,7 @@ def test_dynamicemb_extensions(
             )
             new_size = dyn_emb_rows(table)
             num_evict = evict_counter.cpu().item()
+            # 插入驱逐后的tablesize - 原本的table szize == 新增的值减去驱逐的值 
             assert new_size - old_size == torch.sum(founds == False).item() - num_evict
 
             # convert to torch.int64 to compare, check highest bit not 1.
@@ -292,3 +311,6 @@ def test_dynamicemb_extensions(
         d_num_matched.fill_(0)
         count_matched(table, undump_score, d_num_matched)
         assert d_num_matched.cpu().item() == 0
+
+if __name__ == "__main__":
+    test_dynamicemb_extensions(evict_strategy=EvictStrategy.KLfu, bucket_capacity=128, batch=128, capacity=512 * 1024, num_iteration=8192, dump_interval=1024)
