Check quantized embedding qparam dtypes against how the kernels read them

exir/passes/_quant_patterns_and_replacements.py

@@ -298,7 +298,8 @@ def _(
     num_embeddings, packed_embedding_dim = weight.shape
     embedding_dim = packed_embedding_dim * 4
     embedding = torch.nn.Embedding(num_embeddings, embedding_dim, device=weight.device)
-    return embedding(indices)
+    # As in the reference op, the variant without dtype returns the scales' dtype.
+    return embedding(indices).to(weight_scales.dtype)
 
 
 @register_fake("quantized_decomposed::embedding_2bit.out")
@@ -473,7 +474,8 @@ def _(
     num_embeddings, packed_embedding_dim = weight.shape
     embedding_dim = packed_embedding_dim * 2
     embedding = torch.nn.Embedding(num_embeddings, embedding_dim, device=weight.device)
-    return embedding(indices)
+    # As in the reference op, the variant without dtype returns the scales' dtype.
+    return embedding(indices).to(weight_scales.dtype)
 
 
 @register_fake("quantized_decomposed::embedding_4bit.out")

exir/tests/test_quant_fusion_pass.py

@@ -19,6 +19,7 @@
     QuantFusionPass,
 )
 from executorch.exir.tests.common import register_additional_test_aten_ops
+from torch._subclasses.fake_tensor import FakeTensorMode
 from torch.ao.quantization import (  # @manual
     float_qparams_weight_only_qconfig,
     get_default_qconfig_mapping,
@@ -381,6 +382,34 @@ def forward(self, indices):
             #     m.dump_graph_module().code
             # )
 
+    def test_embedding_xbit_fake_dtype(self) -> None:
+        # The variants without a dtype argument return the scales' dtype; their
+        # fake kernels must report the same dtype, or export records the wrong one.
+        indices = torch.tensor([1, 2, 3])
+        for op, values_per_byte, quant_min, quant_max in (
+            (torch.ops.quantized_decomposed.embedding_2bit.default, 4, -2, 1),
+            (torch.ops.quantized_decomposed.embedding_4bit.default, 2, -8, 7),
+        ):
+            for scales_dtype in (torch.float16, torch.bfloat16, torch.float32):
+                with self.subTest(op=str(op), scales_dtype=scales_dtype):
+                    weight = torch.randint(
+                        0, 256, (10, 32 // values_per_byte), dtype=torch.uint8
+                    )
+                    scales = torch.rand(10, dtype=scales_dtype) + 0.5
+                    expected = op(weight, scales, None, quant_min, quant_max, indices)
+                    with FakeTensorMode() as mode:
+                        actual = op(
+                            mode.from_tensor(weight),
+                            mode.from_tensor(scales),
+                            None,
+                            quant_min,
+                            quant_max,
+                            mode.from_tensor(indices),
+                        )
+                    self.assertEqual(expected.dtype, scales_dtype)
+                    self.assertEqual(actual.dtype, expected.dtype)
+                    self.assertEqual(actual.shape, expected.shape)
+
     def test_embedding_torchao(self) -> None:
         for bit_width, use_dtype_variant, test_per_group, mapping_type in zip(
             [2, 4, 8],

kernels/quantized/cpu/embeddingxb.cpp

@@ -125,9 +125,10 @@ void check_embedding_xbit_args(
         static_cast<int8_t>(weight_scales.dim()));
 
     ET_CHECK_MSG(
-        opt_weight_zero_points.value().scalar_type() == out.scalar_type(),
+        opt_weight_zero_points.value().scalar_type() ==
+            weight_scales.scalar_type(),
         "weight zero points scalar type %" PRId8
-        " does not match out.scalar_type()",
+        " does not match weight_scales.scalar_type()",
         static_cast<int8_t>(opt_weight_zero_points.value().scalar_type()));
 
     for (int32_t i = 0; i < weight_scales.dim(); ++i) {
@@ -290,6 +291,15 @@ Tensor& quantized_embedding_xbit_out(
       out,
       weight_nbit);
 
+  // Without a dtype argument the output takes the scales' dtype, and the
+  // kernel reads the scales as the output type.
+  ET_CHECK_MSG(
+      weight_scales.scalar_type() == out_type,
+      "weight_scales.scalar_type() %" PRId8
+      " does not match out.scalar_type() %" PRId8,
+      static_cast<int8_t>(weight_scales.scalar_type()),
+      static_cast<int8_t>(out_type));
+
   constexpr auto name = "quantized_decomposed::embedding_xbit.out";
   ScalarType indices_type = indices.scalar_type();
   ET_SWITCH_THREE_TYPES(

kernels/quantized/cpu/op_embedding.cpp

@@ -87,9 +87,10 @@ void check_embedding_byte_args(
         static_cast<int8_t>(weight_scales.dim()));
 
     ET_CHECK_MSG(
-        opt_weight_zero_points.value().scalar_type() == out.scalar_type(),
+        opt_weight_zero_points.value().scalar_type() ==
+            weight_scales.scalar_type(),
         "weight zero points scalar type %" PRId8
-        " does not match out.scalar_type()",
+        " does not match weight_scales.scalar_type()",
         static_cast<int8_t>(opt_weight_zero_points.value().scalar_type()));
 
     for (int32_t i = 0; i < weight_scales.dim(); ++i) {
@@ -259,6 +260,15 @@ Tensor& quantized_embedding_byte_out(
       out_type,
       out);
 
+  // Without a dtype argument the output takes the scales' dtype, and the
+  // kernel reads the scales as the output type.
+  ET_CHECK_MSG(
+      weight_scales.scalar_type() == out_type,
+      "weight_scales.scalar_type() %" PRId8
+      " does not match out.scalar_type() %" PRId8,
+      static_cast<int8_t>(weight_scales.scalar_type()),
+      static_cast<int8_t>(out_type));
+
   constexpr auto name = "quantized_decomposed::embedding_byte.out";
   ET_SWITCH_TWO_TYPES(Byte, Char, w_type, ctx, name, CTYPE_W, [&]() {
     ET_SWITCH_THREE_TYPES(

kernels/quantized/test/op_embedding4b_test.cpp

@@ -216,6 +216,79 @@ TEST(OpQuantizedEmbedding4bTest, TestGroupWiseQuantizedEmbeddingBFloat16) {
   EXPECT_TENSOR_CLOSE(out, expected);
 }
 
+TEST(OpQuantizedEmbedding4bTest, DtypeOutHalfQParamsFloatOut) {
+  et_pal_init();
+  TensorFactory<ScalarType::Byte> tfb;
+  TensorFactory<ScalarType::Half> tfh;
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+
+  Tensor weight_scales = tfh.make({3}, {0.5, 1.0, 1.5});
+  Tensor weight_zero_points = tfh.make({3}, {1, -5, 0});
+  Tensor qweight = tfb.make({3, 2}, {89, 239, 163, 72, 11, 126});
+  Tensor indices = tfl.make({3}, {0, 2, 1});
+
+  Tensor out = tf.zeros({3, 4});
+  Tensor expected = tf.make(
+      {3, 4}, {-2.0, 0.0, 2.5, 3.0, -12.0, 4.5, -1.5, 9.0, 7.0, 0.0, 1.0, 5.0});
+
+  quantized_embedding_4bit_dtype_out(
+      qweight,
+      weight_scales,
+      weight_zero_points,
+      -8,
+      7,
+      indices,
+      ScalarType::Float,
+      out);
+
+  EXPECT_TENSOR_CLOSE(out, expected);
+}
+
+TEST(OpQuantizedEmbedding4bTest, DtypeOutZeroPointsScalesMismatchDies) {
+  et_pal_init();
+  TensorFactory<ScalarType::Byte> tfb;
+  TensorFactory<ScalarType::Half> tfh;
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+
+  Tensor weight_scales = tfh.make({3}, {0.5, 1.0, 1.5});
+  Tensor weight_zero_points = tf.make({3}, {1, -5, 0});
+  Tensor qweight = tfb.make({3, 2}, {89, 239, 163, 72, 11, 126});
+  Tensor indices = tfl.make({3}, {0, 2, 1});
+
+  Tensor out = tf.zeros({3, 4});
+  ET_EXPECT_DEATH(
+      quantized_embedding_4bit_dtype_out(
+          qweight,
+          weight_scales,
+          weight_zero_points,
+          -8,
+          7,
+          indices,
+          ScalarType::Float,
+          out),
+      "does not match weight_scales");
+}
+
+TEST(OpQuantizedEmbedding4bTest, OutScalesOutMismatchDies) {
+  et_pal_init();
+  TensorFactory<ScalarType::Byte> tfb;
+  TensorFactory<ScalarType::Half> tfh;
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+
+  Tensor weight_scales = tfh.make({3}, {0.5, 1.0, 1.5});
+  Tensor qweight = tfb.make({3, 2}, {89, 239, 163, 72, 11, 126});
+  Tensor indices = tfl.make({3}, {0, 2, 1});
+
+  Tensor out = tf.zeros({3, 4});
+  ET_EXPECT_DEATH(
+      quantized_embedding_4bit_out(
+          qweight, weight_scales, std::nullopt, -8, 7, indices, out),
+      "does not match out.scalar_type");
+}
+
 TEST(OpQuantizedEmbedding4bTest, TestGroupWiseQuantizedEmbeddingDeath2) {
   et_pal_init();
   TensorFactory<ScalarType::Byte> tfb;

kernels/quantized/test/op_embedding_test.cpp

@@ -510,3 +510,78 @@ TEST(OpQuantizedEmbeddingTest, BFloat16Rounding) {
 
   EXPECT_TENSOR_CLOSE_WITH_TOL(out, expected, 1e-2, 1e-2);
 }
+
+// Zero points share the scales' dtype, which may differ from the output dtype.
+TEST(OpQuantizedEmbeddingTest, DtypeOutHalfQParamsFloatOut) {
+  et_pal_init();
+  TensorFactory<ScalarType::Half> tfh;
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Byte> tfb;
+  TensorFactory<ScalarType::Long> tf_l;
+
+  Tensor weight_scales = tfh.full({3}, 0.5);
+  Tensor weight_zero_points = tfh.full({3}, 1);
+  Tensor qweight = tfb.make({3, 2}, {8, 5, 9, 3, 12, 27});
+  Tensor indices = tf_l.make({2}, {0, 2});
+
+  Tensor out = tf.zeros({2, 2});
+  Tensor expected = tf.make({2, 2}, {3.5, 2, 5.5, 13});
+
+  quantized_embedding_byte_dtype_out(
+      qweight,
+      weight_scales,
+      weight_zero_points,
+      0,
+      255,
+      indices,
+      ScalarType::Float,
+      out);
+
+  EXPECT_TENSOR_CLOSE(out, expected);
+}
+
+TEST(OpQuantizedEmbeddingTest, DtypeOutZeroPointsScalesMismatchDies) {
+  et_pal_init();
+  TensorFactory<ScalarType::Half> tfh;
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Byte> tfb;
+  TensorFactory<ScalarType::Long> tf_l;
+
+  Tensor weight_scales = tfh.full({3}, 0.5);
+  Tensor weight_zero_points = tf.full({3}, 1);
+  Tensor qweight = tfb.make({3, 2}, {8, 5, 9, 3, 12, 27});
+  Tensor indices = tf_l.make({2}, {0, 2});
+
+  Tensor out = tf.zeros({2, 2});
+  ET_EXPECT_DEATH(
+      quantized_embedding_byte_dtype_out(
+          qweight,
+          weight_scales,
+          weight_zero_points,
+          0,
+          255,
+          indices,
+          ScalarType::Float,
+          out),
+      "does not match weight_scales");
+}
+
+// Without a dtype argument the scales are read as the output type, so a
+// narrower scales tensor would be read past its end.
+TEST(OpQuantizedEmbeddingTest, OutScalesOutMismatchDies) {
+  et_pal_init();
+  TensorFactory<ScalarType::Half> tfh;
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Byte> tfb;
+  TensorFactory<ScalarType::Long> tf_l;
+
+  Tensor weight_scales = tfh.full({3}, 0.5);
+  Tensor qweight = tfb.make({3, 2}, {8, 5, 9, 3, 12, 27});
+  Tensor indices = tf_l.make({2}, {0, 2});
+
+  Tensor out = tf.zeros({2, 2});
+  ET_EXPECT_DEATH(
+      quantized_embedding_byte_out(
+          qweight, weight_scales, std::nullopt, 0, 255, indices, out),
+      "does not match out.scalar_type");
+}