diff --git a/include/fbgemm/FbgemmEmbedding.h b/include/fbgemm/FbgemmEmbedding.h
index 835afec943..3caf531976 100644
--- a/include/fbgemm/FbgemmEmbedding.h
+++ b/include/fbgemm/FbgemmEmbedding.h
@@ -51,21 +51,7 @@ FBGEMM_API bool EmbeddingSpMDM(
     bool is_weight_positional = false);
 
 template <typename IndexType>
-class EmbeddingSpMDM4BitKernelSignature {
- public:
-  using Type = std::function<bool(
-      std::int64_t output_size,
-      std::int64_t index_size,
-      std::int64_t data_size,
-      const std::uint8_t* input,
-      const IndexType* indices,
-      const int* lengths,
-      const float* weights, // optional, can be null for non-weighted sum
-      float* out)>;
-};
-
-template <typename IndexType>
-FBGEMM_API typename EmbeddingSpMDM4BitKernelSignature<IndexType>::Type
+FBGEMM_API typename EmbeddingSpMDMKernelSignature<std::uint8_t, IndexType>::Type
 GenerateEmbeddingSpMDM4Bit(
     const std::int64_t block_size,
     bool has_weight,
diff --git a/src/EmbeddingSpMDM4Bit.cc b/src/EmbeddingSpMDM4Bit.cc
index 1130383bba..531e4f7466 100644
--- a/src/EmbeddingSpMDM4Bit.cc
+++ b/src/EmbeddingSpMDM4Bit.cc
@@ -570,7 +570,7 @@ GenEmbeddingSpMDM4BitLookup<indxType>::getOrCreate(
 } // namespace
 
 template <typename indxType>
-typename EmbeddingSpMDM4BitKernelSignature<indxType>::Type
+typename EmbeddingSpMDMKernelSignature<std::uint8_t, indxType>::Type
 GenerateEmbeddingSpMDM4Bit(
     const std::int64_t block_size,
     bool has_weight,
@@ -648,21 +648,23 @@ bool EmbeddingSpMDM4Bit(
       out);
 }
 
-template typename EmbeddingSpMDM4BitKernelSignature<std::int64_t>::Type
-GenerateEmbeddingSpMDM4Bit<std::int64_t>(
-    const std::int64_t block_size,
-    bool has_weight,
-    bool normalize_by_lengths,
-    int prefetch,
-    bool is_weight_positional);
-
-template typename EmbeddingSpMDM4BitKernelSignature<std::int32_t>::Type
-GenerateEmbeddingSpMDM4Bit<std::int32_t>(
-    const std::int64_t block_size,
-    bool has_weight,
-    bool normalize_by_lengths,
-    int prefetch,
-    bool is_weight_positional);
+template
+    typename EmbeddingSpMDMKernelSignature<std::uint8_t, std::int64_t>::Type
+    GenerateEmbeddingSpMDM4Bit<std::int64_t>(
+        const std::int64_t block_size,
+        bool has_weight,
+        bool normalize_by_lengths,
+        int prefetch,
+        bool is_weight_positional);
+
+template
+    typename EmbeddingSpMDMKernelSignature<std::uint8_t, std::int32_t>::Type
+    GenerateEmbeddingSpMDM4Bit<std::int32_t>(
+        const std::int64_t block_size,
+        bool has_weight,
+        bool normalize_by_lengths,
+        int prefetch,
+        bool is_weight_positional);
 
 template bool EmbeddingSpMDM4Bit(
     const std::int64_t block_size,