Add RawIdTrackerWrapper within TBE to access tracked ids and raw ids

Summary:
This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP.

internal 
This is needed to support MPZCH modules for Raw embedding streaming. 
Mode details : 
https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld

Differential Revision: D84925177

Author: aliafzal

torchrec/distributed/batched_embedding_kernel.py

@@ -17,6 +17,7 @@
 from math import sqrt
 from typing import (
     Any,
+    Callable,
     cast,
     Dict,
     Generic,
@@ -70,6 +71,7 @@
     GroupedEmbeddingConfig,
     ShardedEmbeddingTable,
 )
+from torchrec.distributed.model_tracker.types import IndexedLookup
 from torchrec.distributed.shards_wrapper import LocalShardsWrapper
 from torchrec.distributed.types import (
     Shard,
@@ -80,6 +82,7 @@
     TensorProperties,
 )
 from torchrec.distributed.utils import append_prefix, none_throws
+
 from torchrec.modules.embedding_configs import (
     CountBasedEvictionPolicy,
     CountTimestampMixedEvictionPolicy,
@@ -97,13 +100,30 @@
 )
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 RES_ENABLED_TABLES_STR = "res_enabled_tables"
 RES_STORE_SHARDS_STR = "res_store_shards"
 ENABLE_RAW_EMBEDDING_STREAMING_STR = "enable_raw_embedding_streaming"
 
 
+class RawIdTrackerWrapper:
+    def __init__(
+        self,
+        get_indexed_lookups: Callable[
+            [List[str], Optional[str]],
+            List[torch.Tensor],
+        ],
+        delete: Callable[
+            [int],
+            None,
+        ],
+    ) -> None:
+        self.get_indexed_lookups = get_indexed_lookups
+        self.delete = delete
+
+
 def _populate_res_params(config: GroupedEmbeddingConfig) -> Tuple[bool, RESParams]:
     # populate res_params, which is used for raw embedding streaming
     # here only populates the params available in fused_params and TBE configs
@@ -2526,6 +2546,7 @@ def __init__(
         self._lengths_per_emb: List[int] = []
         self.table_name_to_count: Dict[str, int] = {}
         self._param_per_table: Dict[str, TableBatchedEmbeddingSlice] = {}
+        self._raw_id_tracker_wrapper: Optional[RawIdTrackerWrapper] = None
 
         for idx, table_config in enumerate(self._config.embedding_tables):
             self._local_rows.append(table_config.local_rows)
@@ -2579,7 +2600,26 @@ def init_parameters(self) -> None:
                     weight_init_max,
                 )
 
-    def forward(self, features: KeyedJaggedTensor) -> torch.Tensor:
+    def forward(
+        self,
+        features: KeyedJaggedTensor,
+    ) -> torch.Tensor:
+        forward_args: Dict[str, Any] = {}
+        if self._raw_id_tracker_wrapper is not None:
+            if isinstance(self.emb_module, SplitTableBatchedEmbeddingBagsCodegen):
+                raw_id_tracker_wrapper = self._raw_id_tracker_wrapper
+                assert (
+                    raw_id_tracker_wrapper is not None
+                ), "self._raw_id_tracker_wrapper should not be None"
+                # TODO: Calling get_indexed_lookups(None) retrieves raw IDs for ALL tracked FQNs,
+                # including those this TBE doesn't own, and advances the shared consumer read index.
+                # While storage isn't deleted, advancing the index prevents re-reading, which blocks
+                # other TBEs from accessing their tracked raw IDs.
+                raw_ids_list = raw_id_tracker_wrapper.get_indexed_lookups(
+                    features.keys(), self.emb_module.uuid
+                )
+                if raw_ids_list:
+                    forward_args["hash_zch_identities"] = torch.cat(raw_ids_list)
         weights = features.weights_or_none()
         if weights is not None and not torch.is_floating_point(weights):
             weights = None
@@ -2591,17 +2631,22 @@ def forward(self, features: KeyedJaggedTensor) -> torch.Tensor:
                 SSDTableBatchedEmbeddingBags,
             ),
         ):
+            forward_args["batch_size_per_feature_per_rank"] = (
+                features.stride_per_key_per_rank()
+            )
+
+        if len(forward_args) == 0:
             return self.emb_module(
                 indices=features.values().long(),
                 offsets=features.offsets().long(),
                 per_sample_weights=weights,
-                batch_size_per_feature_per_rank=features.stride_per_key_per_rank(),
             )
         else:
             return self.emb_module(
                 indices=features.values().long(),
                 offsets=features.offsets().long(),
                 per_sample_weights=weights,
+                **forward_args,
             )
 
     # pyre-fixme[14]: `state_dict` overrides method defined in `Module` inconsistently.
@@ -2668,6 +2713,22 @@ def named_parameters_by_table(
         for name, param in self._param_per_table.items():
             yield name, param
 
+    def init_raw_id_tracker(
+        self,
+        get_indexed_lookups: Callable[
+            [List[str], Optional[str]],
+            List[torch.Tensor],
+        ],
+        delete: Callable[
+            [int],
+            None,
+        ],
+    ) -> None:
+        if isinstance(self._emb_module, SplitTableBatchedEmbeddingBagsCodegen):
+            self._raw_id_tracker_wrapper = RawIdTrackerWrapper(
+                get_indexed_lookups, delete
+            )
+
 
 class KeyValueEmbeddingBag(BaseBatchedEmbeddingBag[torch.Tensor], FusedOptimizerModule):
     def __init__(