From f73ad797ecd6ec01e35e401e4f3889926ffdf4a0 Mon Sep 17 00:00:00 2001 From: Sankur Kundu Date: Thu, 15 Oct 2026 17:47:36 +0000 Subject: [PATCH 1/2] Check quantized embedding qparam dtypes against how the kernels read them The reference ops (exir/passes/_quant_patterns_and_replacements.py) require weight_zero_points to have the same dtype as weight_scales, and the variants without a dtype argument return the scales' dtype. The kernels instead read the zero points as the scales' type but checked them against the output dtype, and the .out variants read the scales as the output type without checking it. As a result, embedding_byte/2bit/4bit.dtype_out aborted on fp16 scales and zero points with a fp32 output, which the reference accepts, and accepted fp32 zero points with fp16 scales, reading them as fp16. The .out variants with fp16 scales and a fp32 output read past the end of the scales. Check the zero points against the scales, and in the .out variants check the scales against the output. The new tests for the 8-bit and 4-bit ops (2-bit shares the 4-bit code) fail without the change. --- kernels/quantized/cpu/embeddingxb.cpp | 14 +++- kernels/quantized/cpu/op_embedding.cpp | 14 +++- .../quantized/test/op_embedding4b_test.cpp | 73 ++++++++++++++++++ kernels/quantized/test/op_embedding_test.cpp | 75 +++++++++++++++++++ 4 files changed, 172 insertions(+), 4 deletions(-) diff --git a/kernels/quantized/cpu/embeddingxb.cpp b/kernels/quantized/cpu/embeddingxb.cpp index ed5bc3f4397..e3269777fb7 100644 --- a/kernels/quantized/cpu/embeddingxb.cpp +++ b/kernels/quantized/cpu/embeddingxb.cpp @@ -125,9 +125,10 @@ void check_embedding_xbit_args( static_cast(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(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(weight_scales.scalar_type()), + static_cast(out_type)); + constexpr auto name = "quantized_decomposed::embedding_xbit.out"; ScalarType indices_type = indices.scalar_type(); ET_SWITCH_THREE_TYPES( diff --git a/kernels/quantized/cpu/op_embedding.cpp b/kernels/quantized/cpu/op_embedding.cpp index 682e2046c2e..b57d6612d30 100644 --- a/kernels/quantized/cpu/op_embedding.cpp +++ b/kernels/quantized/cpu/op_embedding.cpp @@ -87,9 +87,10 @@ void check_embedding_byte_args( static_cast(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(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(weight_scales.scalar_type()), + static_cast(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( diff --git a/kernels/quantized/test/op_embedding4b_test.cpp b/kernels/quantized/test/op_embedding4b_test.cpp index f9636e14376..95ffb2b7f2d 100644 --- a/kernels/quantized/test/op_embedding4b_test.cpp +++ b/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 tfb; + TensorFactory tfh; + TensorFactory tf; + TensorFactory 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 tfb; + TensorFactory tfh; + TensorFactory tf; + TensorFactory 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 tfb; + TensorFactory tfh; + TensorFactory tf; + TensorFactory 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 tfb; diff --git a/kernels/quantized/test/op_embedding_test.cpp b/kernels/quantized/test/op_embedding_test.cpp index ad27daf866e..c499ebfface 100644 --- a/kernels/quantized/test/op_embedding_test.cpp +++ b/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 tfh; + TensorFactory tf; + TensorFactory tfb; + TensorFactory 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 tfh; + TensorFactory tf; + TensorFactory tfb; + TensorFactory 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 tfh; + TensorFactory tf; + TensorFactory tfb; + TensorFactory 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"); +} From e2978411380ac87ce8d894d2c3e98ceb3e2aeda2 Mon Sep 17 00:00:00 2001 From: Sankur Kundu Date: Wed, 28 Oct 2026 03:09:54 +0000 Subject: [PATCH 2/2] Return the scales' dtype from the 2-bit and 4-bit embedding fakes The fake kernels of embedding_2bit and embedding_4bit (the variants without a dtype argument) returned torch.nn.Embedding(...)(indices), which is always fp32. The reference ops return weight_scales.dtype, so with fp16 or bf16 scales export recorded fp32 for a tensor that is fp16 or bf16 at run time. Return weight_scales.dtype from both fakes. The new test compares each fake with its reference op for fp16, bf16 and fp32 scales; the fp16 and bf16 cases fail without the change. --- .../_quant_patterns_and_replacements.py | 6 ++-- exir/tests/test_quant_fusion_pass.py | 29 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/exir/passes/_quant_patterns_and_replacements.py b/exir/passes/_quant_patterns_and_replacements.py index 463abc37c5e..fceedbe0274 100644 --- a/exir/passes/_quant_patterns_and_replacements.py +++ b/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") diff --git a/exir/tests/test_quant_fusion_pass.py b/exir/tests/test_quant_fusion_pass.py index 499419a1e10..80c2637adfb 100644 --- a/exir/tests/test_quant_fusion_pass.py +++ b/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],