enable KEY_VALUE kernel in ManagedCollisionEmbeddingCollectionSharder (#3308)

jiayulu · facebook-github-bot · commit af6f7bff0eda · 2025-08-21T15:06:56.000-07:00
Summary: Pull Request resolved: #3308 add `KEY_VALUE` kernel in `ManagedCollisionEmbeddingCollectionSharder` Reviewed By: iamzainhuda Differential Revision: D80717266 fbshipit-source-id: 5fb405bf991dc3935ab313c9865463696cb2b1a4
diff --git 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]:

Original file line number	Diff line number	Diff line change
`@@ -303,6 +303,7 @@ def compute_kernels(`
`303`	`303`	`EmbeddingComputeKernel.FUSED.value,`
`304`	`304`	`EmbeddingComputeKernel.FUSED_UVM_CACHING.value,`
`305`	`305`	`EmbeddingComputeKernel.FUSED_UVM.value,`
	`306`	`+ EmbeddingComputeKernel.KEY_VALUE.value,`
`306`	`307`	`]`
`307`	`308`
`308`	`309`	`def sharding_types(self, compute_device_type: str) -> List[str]:`