Create an interface to enable eviction policy (TREC part) (#2481)

Summary:
Pull Request resolved: #2481

To support various types of eviction policy, the HashZchManagedCollisionModule needs to be able to calculate a score (e.g., TTL) for each incoming ID and pass it to the kernel. The latter will make an informed eviction decision based on the existing score associated with the identity (stored in metadata), and a reference value (e.g., current timestamp).

Reviewed By: dracifer, shruthign

Differential Revision: D64163927

fbshipit-source-id: b5e65fbf0568e7a4950988128fd19c1ffa4426ca

Author: Ziliang Zhao

torchrec/distributed/mc_modules.py

@@ -646,6 +646,8 @@ def compute(
                         table: JaggedTensor(
                             values=kjt.values(),
                             lengths=kjt.lengths(),
+                            # TODO: improve this temp solution by passing real weights
+                            weights=torch.tensor(kjt.length_per_key()),
                         )
                     }
                     mcm = self._managed_collision_modules[table]
@@ -660,6 +662,8 @@ def compute(
                     table: JaggedTensor(
                         values=features.values(),
                         lengths=features.lengths(),
+                        # TODO: improve this temp solution by passing real weights
+                        weights=torch.tensor(kjt.length_per_key()),
                     )
                 }
                 mcm = self._managed_collision_modules[table]
@@ -673,6 +677,7 @@ def compute(
                     keys=fns,
                     values=values,
                     lengths=features.lengths(),
+                    # original weights instead of features splits
                     weights=features.weights_or_none(),
                 )
             )

torchrec/modules/mc_modules.py

@@ -76,6 +76,11 @@ def _mcc_lazy_init(
     return (features, created_feature_order, features_order)
 
 
+@torch.fx.wrap
+def _get_length_per_key(kjt: KeyedJaggedTensor) -> torch.Tensor:
+    return torch.tensor(kjt.length_per_key())
+
+
 @torch.no_grad()
 def dynamic_threshold_filter(
     id_counts: torch.Tensor,
@@ -368,6 +373,7 @@ def forward(
                 table: JaggedTensor(
                     values=kjt.values(),
                     lengths=kjt.lengths(),
+                    weights=_get_length_per_key(kjt),
                 )
             }
             mc_input = mc_module(mc_input)