Add util func for kvzch eviction mask

Summary: Adding a util func for kvzch to get a eviction mask using inference threshold. This is used in publish.

Reviewed By: yixin94

Differential Revision: D79045178

fbshipit-source-id: 862488ac6d3de3e01c38a8e89920cddc083f9277

Author: EddyLXJ

torchrec/modules/embedding_configs.py

@@ -194,6 +194,9 @@ class CountBasedEvictionPolicy(VirtualTableEvictionPolicy):
         15  # eviction threshold for count based eviction policy. 0 means no eviction
     )
     decay_rate: float = 0.99  # default decay by default
+    inference_eviction_threshold: int = (
+        eviction_threshold  # eviction threshold for inference count based eviction policy. 0 means no eviction
+    )
 
 
 @dataclass
@@ -203,6 +206,7 @@ class TimestampBasedEvictionPolicy(VirtualTableEvictionPolicy):
     """
 
     eviction_ttl_mins: int = 24 * 60  # 1 day. 0 means no eviction
+    inference_eviction_ttl_mins: int = eviction_ttl_mins  # 0 means no eviction
 
 
 @dataclass
@@ -216,6 +220,13 @@ class CountTimestampMixedEvictionPolicy(VirtualTableEvictionPolicy):
     )
     decay_rate: float = 0.99  # default decay by default
     eviction_ttl_mins: int = 24 * 60  # 1 day. 0 means no eviction based on timestamp
+    inference_eviction_threshold: int = (
+        eviction_threshold  # eviction threshold for inference count based eviction policy. 0 means no eviction based on count
+    )
+
+    inference_eviction_ttl_mins: int = (
+        eviction_ttl_mins  # 0 means no eviction based on timestamp
+    )
 
 
 @dataclass
@@ -227,6 +238,7 @@ class FeatureL2NormBasedEvictionPolicy(VirtualTableEvictionPolicy):
     eviction_threshold: float = (
         0.0  # eviction threshold for feature l2 norm based eviction policy. 0.0 means no eviction
     )
+    inference_eviction_threshold: float = eviction_threshold
 
 
 @dataclass