CodeLinaro
diff --git a/‎backends/native/serialization/graph_serialize.py‎
Lines changed: 18 additions & 15 deletions b/‎backends/native/serialization/graph_serialize.py‎
Lines changed: 18 additions & 15 deletions
diff --git a/‎backends/native/test/test_quant.py‎
Lines changed: 33 additions & 4 deletions b/‎backends/native/test/test_quant.py‎
Lines changed: 33 additions & 4 deletions
@@ -1054,15 +1054,18 @@ def _fold_gguf_dequant(graph_module: torch.fx.GraphModule) -> None:
         node.meta["native_skip_serialize"] = True
 
 
-def _is_linear_weight_use(dequantize: torch.fx.Node, user: torch.fx.Node) -> bool:
+# The argument through which each op reads a foldable weight.
+_WEIGHT_ARG_INDEX = {"aten::linear": 1, "aten::embedding": 0}
+
+
+def _is_weight_use(dequantize: torch.fx.Node, user: torch.fx.Node) -> bool:
     op = _resolve_op_overload(user.target) if user.op == "call_function" else None
+    index = _WEIGHT_ARG_INDEX.get(op.name()) if op is not None else None
     return (
-        op is not None
-        and op.name() == "aten::linear"
-        and len(user.args) >= 2
-        and user.args[1] is dequantize
-        and user.args[0] is not dequantize
-        and dequantize not in user.args[2:]
+        index is not None
+        and len(user.args) > index
+        and user.args[index] is dequantize
+        and sum(arg is dequantize for arg in user.args) == 1
         and dequantize not in user.kwargs.values()
     )
 
@@ -1072,13 +1075,13 @@ def _mark_torchao_q4_weights(
 ) -> None:
     """Mark portable torchao q4 weights for packed PTN storage.
 
-    A weight dequantize read only as the weight of `aten.linear`, whose output
-    dtype matches its scales and whose operands are serialized state inputs, is
-    folded: the linear reads the packed `AffineGroup` weight directly and the
-    dequantize is not serialized. A weight is packed only if every one of its
-    readers is such a dequantize with the same parameters. Any other dequantize
-    stays in the graph over the plain int8 weight, with torchao semantics.
-    Activation q/dq nodes are never folded.
+    A weight dequantize read only as the weight of `aten.linear` or
+    `aten.embedding`, whose output dtype matches its scales and whose operands
+    are serialized state inputs, is folded: the op reads the packed `AffineGroup`
+    weight directly and the dequantize is not serialized. A weight is packed only
+    if every one of its readers is such a dequantize with the same parameters.
+    Any other dequantize stays in the graph over the plain int8 weight, with
+    torchao semantics. Activation q/dq nodes are never folded.
     """
     if lifted_constants is None:
         return
@@ -1132,7 +1135,7 @@ def _mark_torchao_q4_weights(
             {weight.name, scale.name, zero_point.name} <= lifted_constants
             and output.dtype == scale_value.dtype
             and node.users
-            and all(_is_linear_weight_use(node, user) for user in node.users)
+            and all(_is_weight_use(node, user) for user in node.users)
         ):
             folds.setdefault(weight, []).append((node, packed))
 
 
@@ -128,9 +128,9 @@ def test_rejects_values_outside_signed_int4_range(self):
                     _pack_signed_int4(torch.tensor([[value, 0]], dtype=torch.int8))
 
 
-def _linear_weight(method):
-    [linear] = [n for n in method.graph.nodes if n.target and "aten.linear" in n.target]
-    [weight] = [a.arg.value for a in linear.inputs if a.name == "weight"]
+def _weight_of(method, op="aten.linear"):
+    [node] = [n for n in method.graph.nodes if n.target and op in n.target]
+    [weight] = [a.arg.value for a in node.inputs if a.name == "weight"]
     assert isinstance(weight, TensorArg)
     return weight.name
 
@@ -148,7 +148,7 @@ def _assert_linear_reads_packed_weight(self, method, constants):
         [weight] = [c for c in method.constants if c.meta.quant is not None]
         scheme = weight.meta.quant.scheme
         self.assertIsInstance(scheme, AffineGroup)
-        self.assertEqual(_linear_weight(method), weight.name)
+        self.assertEqual(_weight_of(method), weight.name)
         self.assertEqual(weight.meta.dtype, ScalarType.BYTE)
         self.assertEqual([d.max for d in weight.meta.sizes], [self.N, self.K])
         self.assertEqual(
@@ -185,6 +185,35 @@ def test_4w_linear_reads_packed_weight(self):
         self._assert_linear_reads_packed_weight(method, constants)
 
 
+class TorchaoQ4EmbeddingTest(unittest.TestCase):
+    def test_embedding_reads_packed_weight(self):
+        rows, cols, group_size = 16, 64, 32
+        model = nn.Sequential(nn.Embedding(rows, cols))
+        quantize_(
+            model,
+            IntxWeightOnlyConfig(
+                weight_dtype=torch.int4, granularity=PerGroup(group_size)
+            ),
+            filter_fn=lambda m, _: isinstance(m, nn.Embedding),
+        )
+        indices = torch.tensor([[1, 5, 7]])
+        manager = to_native(torch.export.export(model, (indices,)))
+        method = deserialize_program(manager._ptg).methods[0]
+
+        self.assertFalse(
+            any("dequantize_affine" in t for t in _call_function_targets(method.graph))
+        )
+        [weight] = [c for c in method.constants if c.meta.quant is not None]
+        scheme = weight.meta.quant.scheme
+        self.assertEqual(_weight_of(method, "aten.embedding"), weight.name)
+        self.assertEqual([d.max for d in weight.meta.sizes], [rows, cols])
+        self.assertEqual(
+            (scheme.quant_min, scheme.quant_max, scheme.group_size),
+            (-8, 7, group_size),
+        )
+        self.assertEqual(manager._constants[weight.data_key].numel(), rows * cols // 2)
+
+
 class _DequantizedLinear(nn.Module):
     """linear(x, dequantize_affine(weight)) over an int4 [8, 16] weight in groups
     of 8, optionally with the weight as an input or the dequantize read twice."""