Add VBE support to sharded module prefetch (#1809)

Summary:
Pull Request resolved: #1809

enable prefetch for VBE case

Reviewed By: sryap, henrylhtsang

Differential Revision: D54816141

fbshipit-source-id: 6bd5b27b3c1b8f26074f48b83e223ce51115afe5

torchrec/distributed/embedding_lookup.py

@@ -425,6 +425,11 @@ def _need_prefetch(config: GroupedEmbeddingConfig) -> bool:
                         indices=features.values(),
                         offsets=features.offsets(),
                         forward_stream=forward_stream,
+                        batch_size_per_feature_per_rank=(
+                            features.stride_per_key_per_rank()
+                            if features.variable_stride_per_key()
+                            else None
+                        ),
                     )
 
     def forward(