diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py index 4d55ed2738..bd43100cb0 100644 --- a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py +++ b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py @@ -240,6 +240,19 @@ def validate(self) -> None: ), "backend_return_whole_row can only be enabled when enable_optimizer_offloading is enabled" +class KVZCHEvictionTBEConfig(NamedTuple): + # Eviction trigger model for kvzch table: 0: disabled, 1: iteration, 2: mem_util, 3: manual, 4: id count, 5: free_mem + kvzch_eviction_trigger_mode: int = 2 # mem_util + # Minimum free memory (in GB) required before triggering eviction when using free_mem trigger mode. + eviction_free_mem_threshold_gb: int = 200 # 200GB + # Number of batches between checks for free memory threshold when using free_mem trigger mode. + eviction_free_mem_check_interval_batch: int = 1000 + # The width of each feature score bucket used for threshold calculation in feature score-based eviction. + threshold_calculation_bucket_stride: float = 0.2 + # Total number of feature score buckets used for threshold calculation in feature score-based eviction. + threshold_calculation_bucket_num: Optional[int] = 1000000 # 1M + + class BackendType(enum.IntEnum): SSD = 0 DRAM = 1