Add eviction writeback callback to FeatureEvict (pytorch#5832)

Summary:

X-link: facebookresearch/FBGEMM#2754

Add EvictionWritebackCallback type and wire it through FeatureEvict eviction loops. When a dirty block is evicted from DRAM, its data is serialized and passed to the callback for writeback to SSD. The callback is invoked after releasing the DRAM shard lock to avoid holding locks during SSD I/O.

Differential Revision: D107491725

Author: Lizhe Ji

fbgemm_gpu/fbgemm_gpu/tbe/ssd/ssd_config.py

@@ -83,6 +83,9 @@ class EvictionPolicy(NamedTuple):
     enable_eviction_for_feature_score_eviction_policy: list[bool] | None = (
         None  # enable eviction if eviction policy is feature score, false means no eviction
     )
+    enable_ssd_writeback: bool = (
+        False  # when True, evicted dirty blocks are written back to SSD via callback; when False, eviction deletes blocks directly
+    )
 
     def validate(self) -> None:
         assert self.eviction_trigger_mode in [0, 1, 2, 3, 4, 5], (

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -956,6 +956,7 @@ def __init__(
                     eviction_policy.interval_for_insufficient_eviction_s,
                     eviction_policy.interval_for_sufficient_eviction_s,
                     eviction_policy.interval_for_feature_statistics_decay_s,
+                    eviction_policy.enable_ssd_writeback,
                 )
             enrichment_config = None
             if (