Serialize GGUF weights as packed constants (pytorch#22856)

Summary:

Keep GGUF-quantized weights in their packed block form through lowering and
serialization, instead of materializing a dequantized float weight inside the
delegate.

- The partitioner claims torchao::dequantize_gguf (importing
  executorch.extension.llm.export.gguf to register it), so a GGUF weight
  dequantize stays inside the delegate rather than forcing a partition boundary.

- graph_serialize._fold_gguf_dequant folds the dequantize at write time, the way
  XNNPACK folds a weight dequantize. Graph topology is not rewritten; instead
  each dequantize_gguf node is tagged to skip serialization, its packed weight
  placeholder is tagged with the logical shape and codec, and references to the
  dequantize output are redirected to the packed weight. The consuming
  linear/embedding is therefore emitted directly over the packed constant.

- The packed constant serializes as TensorMeta(dtype=BYTE, sizes=<logical
  shape>, quant=QuantSpec(scheme=PackedQuant(codec="gguf:<qtype>"))). The raw
  GGUF block bytes ship as-is, and the runtime dequantizes inside the
  matmul/gather.

Reviewed By: digantdesai

Differential Revision: D111753294

Pull Request resolved: pytorch#22856

Author: metascroy

backends/native/BUCK

@@ -42,5 +42,6 @@ fbcode_target(
         "//executorch/exir/backend:utils",
         "//executorch/exir/passes:cse_pass",
         "//executorch/exir/passes:lib",
+        "//executorch/extension/llm/export:gguf",
     ],
 )

backends/native/partitioner.py

@@ -15,6 +15,8 @@
 
 from typing import Callable, final, List, Mapping, Optional, Tuple
 
+# Registers torch.ops.torchao.dequantize_gguf, referenced in _SUPPORTED_NON_CORE_OPS.
+import executorch.extension.llm.export.gguf  # noqa: F401
 import torch
 from executorch.backends.native.passes import backend_inplace_aten_variants
 
@@ -38,6 +40,9 @@
     torch.ops.aten.linear.default,
     torch.ops.aten.addmm.default,
     torch.ops.aten.scaled_dot_product_attention.default,
+    # GGUF weight dequantize stays in the delegate; the serializer folds it into a
+    # PackedQuant weight on the consuming op.
+    torch.ops.torchao.dequantize_gguf.default,
     torch.ops.aten.rms_norm.default,
 ]
 

backends/native/serialization/graph_serialize.py

@@ -66,7 +66,9 @@
     OutputKind,
     OutputSpec,
     OutputValueKind,
+    PackedQuant,
     Program,
+    QuantSpec,
     ScalarType,
     ScalarTypeArg,
     StringArg,
@@ -229,6 +231,17 @@ def _tensor_meta(t: torch.Tensor) -> TensorMeta:
     )
 
 
+def _packed_tensor_meta(codec: str, sizes: tuple[int, ...]) -> TensorMeta:
+    """TensorMeta for a codec-packed weight: logical shape with BYTE dtype, the
+    packing described by a PackedQuant scheme. The raw block bytes ship as data.
+    """
+    return TensorMeta(
+        dtype=ScalarType.BYTE,
+        sizes=[_dim(int(s)) for s in sizes],
+        quant=QuantSpec(scheme=PackedQuant(codec=codec)),
+    )
+
+
 # ---------------------------------------------------------------------------
 # Argument dispatch.
 # ---------------------------------------------------------------------------
@@ -353,6 +366,11 @@ def _node_to_arg_value(
     v: torch.fx.Node,
     subgraph_map: "dict[str, torch.fx.GraphModule] | None",
 ) -> ArgumentValue:
+    # A GGUF weight consumed through dequantize_gguf serializes as the raw packed
+    # constant (see _fold_gguf_dequant), so redirect the reference to it.
+    redirect = v.meta.get("native_serialize_as")
+    if redirect is not None:
+        return TensorArg(name=redirect)
     # A get_attr node that resolves to a submodule GraphModule is a
     # higher-order-op subgraph (torch.cond branch, map body, ...). Inline it
     # as a GraphArg rather than referencing it like a tensor value.
@@ -758,6 +776,11 @@ def _collect_fx_nodes(
             ):
                 continue
 
+        # Folded GGUF dequantize nodes are not emitted; the consuming op references
+        # the packed weight directly (see _fold_gguf_dequant).
+        if fx_node.meta.get("native_skip_serialize"):
+            continue
+
         kind = _op_kind(fx_node.op)
         target = None
         inputs: list[NamedArgument] = []
@@ -839,10 +862,12 @@ def _extract_constants_and_mutable_buffers(
     state_dict: dict[str, object],
     constants: dict[str, object] | None,
     mutated_fqns: set[str],
+    packed_quant: dict[str, dict[str, object]] | None = None,
 ) -> tuple[list[NamedTensorRef], dict[str, torch.Tensor], list[MutableBufferSpec]]:
     constant_refs: list[NamedTensorRef] = []
     constant_data: dict[str, torch.Tensor] = {}
     mutable_buffers: list[MutableBufferSpec] = []
+    packed_quant = packed_quant or {}
 
     for ispec in getattr(graph_signature, "input_specs", []) or []:
         if ispec.kind not in _INPUT_KIND_MAP:
@@ -862,11 +887,17 @@ def _extract_constants_and_mutable_buffers(
         if not isinstance(tensor, torch.Tensor):
             continue
         tensor = tensor.contiguous()
+        pq = packed_quant.get(name)
+        meta = (
+            _packed_tensor_meta(pq["codec"], pq["sizes"])
+            if pq is not None
+            else _tensor_meta(tensor)
+        )
         constant_refs.append(
             NamedTensorRef(
                 name=name,
                 data_key=target_fqn,
-                meta=_tensor_meta(tensor),
+                meta=meta,
                 kind=_INPUT_KIND_MAP[ispec.kind],
                 mutated=target_fqn in mutated_fqns,
             )
@@ -916,9 +947,14 @@ def _build_method_graph(
 ]:
     _validate_tensor_user_inputs(graph_signature)
     output_specs, mutated_fqns = _build_output_specs(output_names, graph_signature)
+    packed_quant = {
+        n.name: n.meta["native_packed_quant"]
+        for n in graph_module.graph.nodes
+        if n.op == "placeholder" and "native_packed_quant" in n.meta
+    }
     constant_refs, constant_data, mutable_buffers = (
         _extract_constants_and_mutable_buffers(
-            graph_signature, state_dict, constants, mutated_fqns
+            graph_signature, state_dict, constants, mutated_fqns, packed_quant
         )
     )
     constant_names = {c.name for c in constant_refs}
@@ -936,6 +972,30 @@ def _build_method_graph(
     return graph, constant_refs, output_specs, mutable_buffers, constant_data
 
 
+def _fold_gguf_dequant(graph_module: torch.fx.GraphModule) -> None:
+    """Tag GGUF dequantize nodes so serialization emits the consuming op over the
+    packed weight, the way XNNPACK folds a weight dequantize at write time.
+
+    The graph is left unchanged: each dequantize_gguf node is marked to skip, its
+    packed weight arg is tagged with a PackedQuant scheme (logical shape, codec),
+    and references to the dequantize output are redirected to the packed weight.
+    """
+    for node in graph_module.graph.nodes:
+        op = _resolve_op_overload(node.target) if node.op == "call_function" else None
+        if op is None or op.name() != "torchao::dequantize_gguf" or len(node.args) < 2:
+            continue
+        weight = node.args[0]
+        val = node.meta.get("val")
+        if not isinstance(weight, torch.fx.Node) or not isinstance(val, torch.Tensor):
+            continue
+        weight.meta["native_packed_quant"] = {
+            "codec": f"gguf:{node.args[1]}",
+            "sizes": tuple(int(s) for s in val.shape),
+        }
+        node.meta["native_serialize_as"] = weight.name
+        node.meta["native_skip_serialize"] = True
+
+
 def _build_graph_body(
     graph_module: torch.fx.GraphModule,
     graph_signature: object | None,
@@ -948,6 +1008,7 @@ def _build_graph_body(
     list[MutableBufferSpec],
     dict[str, torch.Tensor],
 ]:
+    _fold_gguf_dequant(graph_module)
     subgraph_map = _subgraph_map(graph_module)
     nodes, val_by_name, output_names = _collect_fx_nodes(graph_module, subgraph_map)
 

backends/native/test/BUCK

@@ -22,6 +22,7 @@ fbcode_target(
         "//caffe2:torch",
         "//executorch/backends/native:lib",
         "//executorch/exir:lib",
+        "//executorch/extension/llm/export:gguf",
     ],
 )
 
@@ -45,6 +46,20 @@ fbcode_target(
         "//caffe2:torch",
         "//executorch/backends/native:lib",
         "//executorch/exir:lib",
+        "//executorch/extension/llm/export:gguf",
+    ],
+)
+
+fbcode_target(
+    _kind = runtime.python_test,
+    name = "test_quant",
+    srcs = ["test_quant.py"],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/native:lib",
+        "//executorch/backends/native/test:utils",
+        "//executorch/exir:lib",
+        "//executorch/extension/llm/export:gguf",
     ],
 )
 
@@ -56,5 +71,6 @@ fbcode_target(
         "//caffe2:torch",
         "//executorch/backends/native:lib",
         "//executorch/exir:lib",
+        "//executorch/extension/llm/export:gguf",
     ],
 )

backends/native/test/test_quant.py

@@ -0,0 +1,98 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+import torch.nn as nn
+
+from executorch.backends.native.serialization import deserialize_program
+from executorch.backends.native.serialization.schema import PackedQuant, ScalarType
+from executorch.backends.native.test.utils import (
+    _call_function_targets,
+    _get_delegate_blob,
+    _lower,
+)
+from executorch.extension.llm.export.gguf import ExportableGGUFTensor
+
+
+class FuseGGUFPassTest(unittest.TestCase):
+    def _packed_constants(self, program):
+        return [
+            c
+            for c in (program.methods[0].constants or [])
+            if c.meta.quant is not None and isinstance(c.meta.quant.scheme, PackedQuant)
+        ]
+
+    def test_gguf_linear_serializes_as_packed_linear(self):
+        # A GGUF-quantized linear serializes as a plain linear over a weight
+        # constant tagged with a PackedQuant codec, dropping the dequantize.
+        n, k = 8, 256  # k must be a multiple of QK_K (256); q4_k block = 144 bytes
+        blob = torch.randint(0, 256, (n, (k // 256) * 144), dtype=torch.uint8)
+        lin = nn.Linear(k, n, bias=False)
+        lin.weight = nn.Parameter(
+            ExportableGGUFTensor.from_raw(blob, "q4_k", torch.float32),
+            requires_grad=False,
+        )
+
+        program = deserialize_program(
+            _get_delegate_blob(_lower(lin, (torch.randn(2, k),)))
+        )
+        targets = _call_function_targets(program.methods[0].graph)
+        self.assertTrue(any(t and "linear" in t for t in targets))
+        self.assertFalse(any(t and "dequantize_gguf" in t for t in targets))
+
+        packed = self._packed_constants(program)
+        self.assertEqual(len(packed), 1)
+        self.assertEqual(packed[0].meta.quant.scheme.codec, "gguf:q4_k")
+        self.assertEqual(packed[0].meta.dtype, ScalarType.BYTE)
+
+    def test_activation_dtype_in_metadata(self):
+        # The activation compute dtype must round-trip into the serialized
+        # TensorMeta of every value in the graph (the packed weight stays BYTE).
+        n, k = 8, 256  # k multiple of QK_K (256); q4_k block = 144 bytes
+        blob = torch.randint(0, 256, (n, (k // 256) * 144), dtype=torch.uint8)
+        for dtype, expected in (
+            (torch.float32, ScalarType.FLOAT),
+            (torch.float16, ScalarType.HALF),
+            (torch.bfloat16, ScalarType.BFLOAT16),
+        ):
+            with self.subTest(dtype=dtype):
+                lin = nn.Linear(k, n, bias=False)
+                lin.weight = nn.Parameter(
+                    ExportableGGUFTensor.from_raw(blob, "q4_k", dtype),
+                    requires_grad=False,
+                )
+                program = deserialize_program(
+                    _get_delegate_blob(_lower(lin, (torch.randn(2, k, dtype=dtype),)))
+                )
+                metas = [
+                    tv.meta for tv in (program.methods[0].graph.tensor_values or [])
+                ]
+                self.assertTrue(metas, "expected serialized tensor values")
+                for meta in metas:
+                    self.assertEqual(meta.dtype, expected)
+
+    def test_gguf_embedding_serializes_as_packed_embedding(self):
+        # The embedding counterpart: plain embedding over a PackedQuant weight.
+        num_emb, k = 16, 256  # k multiple of QK_K (256); q4_k block = 144 bytes
+        blob = torch.randint(0, 256, (num_emb, (k // 256) * 144), dtype=torch.uint8)
+        emb = nn.Embedding(num_emb, k)
+        emb.weight = nn.Parameter(
+            ExportableGGUFTensor.from_raw(blob, "q4_k", torch.float32),
+            requires_grad=False,
+        )
+
+        program = deserialize_program(
+            _get_delegate_blob(_lower(emb, (torch.randint(0, num_emb, (2, 3)),)))
+        )
+        targets = _call_function_targets(program.methods[0].graph)
+        self.assertTrue(any(t and "embedding" in t for t in targets))
+        self.assertFalse(any(t and "dequantize_gguf" in t for t in targets))
+
+        packed = self._packed_constants(program)
+        self.assertEqual(len(packed), 1)
+        self.assertEqual(packed[0].meta.quant.scheme.codec, "gguf:q4_k")