From ed43af91020ae80c34432d4ba69811f50fb5ccbb Mon Sep 17 00:00:00 2001
From: Yulu Jia <yulujia@meta.com>
Date: Thu, 21 Aug 2025 10:35:10 -0700
Subject: [PATCH] enable KEY_VALUE kernel in
 ManagedCollisionEmbeddingCollectionSharder

Summary: add `KEY_VALUE` kernel in `ManagedCollisionEmbeddingCollectionSharder`

Reviewed By: iamzainhuda

Differential Revision: D80717266
---
 torchrec/distributed/mc_embedding_modules.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/torchrec/distributed/mc_embedding_modules.py b/torchrec/distributed/mc_embedding_modules.py
index b817f020a..63673be73 100644
--- a/torchrec/distributed/mc_embedding_modules.py
+++ b/torchrec/distributed/mc_embedding_modules.py
@@ -303,6 +303,7 @@ def compute_kernels(
             EmbeddingComputeKernel.FUSED.value,
             EmbeddingComputeKernel.FUSED_UVM_CACHING.value,
             EmbeddingComputeKernel.FUSED_UVM.value,
+            EmbeddingComputeKernel.KEY_VALUE.value,
         ]
 
     def sharding_types(self, compute_device_type: str) -> List[str]: