Add uvm caching compute kernel to mch sharders

Summary: Add uvm caching to allowed kernels.

Reviewed By: dstaay-fb

Differential Revision: D52919490

torchrec/distributed/mc_embedding_modules.py

@@ -264,7 +264,10 @@ def compute_kernels(
         sharding_type: str,
         compute_device_type: str,
     ) -> List[str]:
-        return [EmbeddingComputeKernel.FUSED.value]
+        return [
+            EmbeddingComputeKernel.FUSED.value,
+            EmbeddingComputeKernel.FUSED_UVM_CACHING.value,
+        ]
 
     def sharding_types(self, compute_device_type: str) -> List[str]:
         return list(