Skip to content

Commit 70437b3

Browse files
metascroymetascroy
committed
up
1 parent ec7c5c4 commit 70437b3

21 files changed

Lines changed: 713 additions & 359 deletions

File tree

‎.github/workflows/mlx.yml‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ jobs:
7777
backends/mlx/test/test_passes.py \
7878
backends/mlx/test/test_pattern_utils.py \
7979
backends/mlx/test/test_partitioner.py \
80+
backends/mlx/test/test_serialization_dedup.py \
8081
examples/models/gemma4_31b/tests/test_mlx_pipeline.py \
8182
-v
8283
echo "::endgroup::"
@@ -90,11 +91,12 @@ jobs:
9091
echo "::endgroup::"
9192
9293
echo "::group::Run custom_kernel_ops op tests"
93-
# Run every custom_kernel_ops/test/test_*.py via its OpTestCase `run` CLI. Adding a
94+
# Run every custom_kernel_ops/**/test/test_*.py via its OpTestCase `run`
95+
# CLI. Recurses into per-format subpackages (e.g. gguf/test), so adding a
9496
# new op test file requires no change here.
9597
set -e
96-
for t in backends/mlx/custom_kernel_ops/test/test_*.py; do
97-
mod="executorch.backends.mlx.custom_kernel_ops.test.$(basename "$t" .py)"
98+
for t in $(find backends/mlx/custom_kernel_ops -path '*/test/test_*.py' | sort); do
99+
mod="executorch.$(echo "${t%.py}" | tr '/' '.')"
98100
echo "--- ${mod} ---"
99101
${CONDA_RUN} python -m "${mod}" run -v
100102
done

‎backends/mlx/custom_kernel_ops/gguf/__init__.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
#
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
#
8+
9+
"""GGUF-quantized custom ops for the MLX backend.
10+
11+
Submodules:
12+
13+
* :mod:`.q6k` -- shared Q6_K primitives (constants, pure-torch dequant,
14+
Metal header). Import this for symbols; it does not register any op.
15+
* :mod:`.linear` -- registers ``mlx::gguf_linear``.
16+
* :mod:`.embedding` -- registers ``mlx::gguf_embedding``.
17+
18+
To register an op, import the corresponding submodule for its side effect, e.g.
19+
``import executorch.backends.mlx.custom_kernel_ops.gguf.linear # noqa: F401``.
20+
21+
This package ``__init__`` is intentionally side-effect free so importing
22+
``.q6k`` for the pure-torch dequant does not pull in the MLX builder/registry.
23+
"""

‎backends/mlx/custom_kernel_ops/gguf/embedding.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
#
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
#
8+
9+
"""``mlx::gguf_embedding``: embedding gather against a GGUF-quantized table.
10+
11+
out = dequant(weight[indices])
12+
13+
This module is a thin **format router**: it owns the ``mlx::gguf_embedding`` op
14+
identity (custom op, fake, and lowering registration) and dispatches on
15+
``format`` to a per-format implementation. Only ``"q6k"`` is supported today
16+
(see :mod:`.q6k.embedding`); other formats raise ``NotImplementedError``.
17+
18+
Usage::
19+
20+
import executorch.backends.mlx.custom_kernel_ops.gguf.embedding # noqa: F401
21+
22+
out = torch.ops.mlx.gguf_embedding(weight, indices, "q6k")
23+
# weight: (vocab, (K/256)*210) uint8 GGUF q6_K blob
24+
# indices: (...) int
25+
# out: (..., K) bfloat16
26+
"""
27+
28+
from __future__ import annotations
29+
30+
import torch
31+
from torch import Tensor
32+
from torch.fx.node import Node
33+
34+
35+
# ---------------------------------------------------------------------------
36+
# Custom op + eager fallback (format-agnostic shell; dispatches by format)
37+
# ---------------------------------------------------------------------------
38+
39+
40+
@torch.library.custom_op("mlx::gguf_embedding", mutates_args=())
41+
def gguf_embedding(weight: Tensor, indices: Tensor, format: str) -> Tensor:
42+
"""Gather + dequantize rows of a GGUF-quantized embedding table.
43+
44+
Args:
45+
weight: ``(vocab, row_bytes)`` uint8 GGUF quant blob.
46+
indices: integer token ids of any shape.
47+
format: GGUF quant type; only ``"q6k"`` supported.
48+
49+
Returns:
50+
``(*indices.shape, K)`` bfloat16.
51+
"""
52+
if format == "q6k":
53+
from executorch.backends.mlx.custom_kernel_ops.gguf.q6k.embedding import (
54+
eager_embedding,
55+
)
56+
57+
return eager_embedding(weight, indices)
58+
raise NotImplementedError(
59+
f"mlx::gguf_embedding: unsupported format {format!r}; only 'q6k' is supported"
60+
)
61+
62+
63+
@torch.library.register_fake("mlx::gguf_embedding")
64+
def gguf_embedding_fake(weight: Tensor, indices: Tensor, format: str) -> Tensor:
65+
from executorch.backends.mlx.custom_kernel_ops.gguf.q6k.common import (
66+
Q6K_BLOCK_BYTES,
67+
QK_K,
68+
)
69+
70+
row_bytes = weight.shape[1]
71+
K = (row_bytes // Q6K_BLOCK_BYTES) * QK_K
72+
return indices.new_empty((*indices.shape, K), dtype=torch.bfloat16)
73+
74+
75+
# ---------------------------------------------------------------------------
76+
# MLX handler (format router)
77+
# ---------------------------------------------------------------------------
78+
79+
from executorch.backends.mlx.builder.op_registry import REGISTRY
80+
from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder
81+
from executorch.backends.mlx.builder.slot_manager import Slot
82+
83+
84+
@REGISTRY.register(target=[torch.ops.mlx.gguf_embedding.default])
85+
def _gguf_embedding_handler(P: MLXProgramBuilder, n: Node) -> Slot:
86+
"""Route ``mlx::gguf_embedding`` lowering to the per-format implementation."""
87+
args = P.args(n)
88+
fmt = args[2] if len(args) >= 3 else None
89+
if fmt == "q6k":
90+
from executorch.backends.mlx.custom_kernel_ops.gguf.q6k.embedding import (
91+
emit_embedding,
92+
)
93+
94+
return emit_embedding(P, n)
95+
raise NotImplementedError(
96+
f"mlx::gguf_embedding: unsupported format {fmt!r}; only 'q6k' is supported"
97+
)

‎backends/mlx/custom_kernel_ops/gguf/linear.py‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
#
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
#
8+
9+
"""``mlx::gguf_linear``: linear layer against a GGUF-quantized weight.
10+
11+
out = x @ dequant(weight)^T (+ bias)
12+
13+
The weight is stored in the **exact GGUF packed block layout** (no repacking),
14+
so weights converted by llama.cpp / gguf-py can be consumed directly. The
15+
``format`` argument selects the GGUF quantization type.
16+
17+
This module is a thin **format router**: it owns the ``mlx::gguf_linear`` op
18+
identity (custom op, fake, and lowering registration) and dispatches on
19+
``format`` to a per-format implementation. Only ``"q6k"`` is supported today
20+
(see :mod:`.q6k.linear`); other formats raise ``NotImplementedError``. To add a
21+
format, implement ``eager_linear`` / ``emit_linear`` in a sibling package (e.g.
22+
``q4k``) and add a branch below.
23+
24+
Usage::
25+
26+
import executorch.backends.mlx.custom_kernel_ops.gguf.linear # noqa: F401
27+
28+
out = torch.ops.mlx.gguf_linear(x, weight, "q6k", bias)
29+
# x: (..., K) bf16 / fp16 / fp32
30+
# weight: (N, (K/256)*210) uint8 GGUF q6_K blob
31+
# bias: (N,) or None activation dtype
32+
# out: (..., N) activation dtype
33+
"""
34+
35+
from __future__ import annotations
36+
37+
from typing import Optional
38+
39+
import torch
40+
from torch import Tensor
41+
from torch.fx.node import Node
42+
43+
44+
# ---------------------------------------------------------------------------
45+
# Custom op + eager fallback (format-agnostic shell; dispatches by format)
46+
# ---------------------------------------------------------------------------
47+
48+
49+
@torch.library.custom_op("mlx::gguf_linear", mutates_args=())
50+
def gguf_linear(
51+
x: Tensor,
52+
weight: Tensor,
53+
format: str,
54+
bias: Optional[Tensor] = None,
55+
) -> Tensor:
56+
"""Linear against a GGUF-quantized weight.
57+
58+
Args:
59+
x: ``(..., K)`` activations (bf16 / fp16 / fp32).
60+
weight: ``(N, row_bytes)`` uint8 GGUF quant blob.
61+
format: GGUF quant type; only ``"q6k"`` supported.
62+
bias: optional ``(N,)`` of activation dtype.
63+
64+
Returns:
65+
``(..., N)`` of activation dtype.
66+
"""
67+
if format == "q6k":
68+
from executorch.backends.mlx.custom_kernel_ops.gguf.q6k.linear import (
69+
eager_linear,
70+
)
71+
72+
return eager_linear(x, weight, bias)
73+
raise NotImplementedError(
74+
f"mlx::gguf_linear: unsupported format {format!r}; only 'q6k' is supported"
75+
)
76+
77+
78+
@torch.library.register_fake("mlx::gguf_linear")
79+
def gguf_linear_fake(
80+
x: Tensor,
81+
weight: Tensor,
82+
format: str,
83+
bias: Optional[Tensor] = None,
84+
) -> Tensor:
85+
N = weight.shape[0]
86+
out_shape = list(x.shape)
87+
out_shape[-1] = N
88+
return x.new_empty(out_shape, dtype=x.dtype)
89+
90+
91+
# ---------------------------------------------------------------------------
92+
# MLX handler (format router)
93+
# ---------------------------------------------------------------------------
94+
95+
from executorch.backends.mlx.builder.op_registry import REGISTRY
96+
from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder
97+
from executorch.backends.mlx.builder.slot_manager import Slot
98+
99+
100+
@REGISTRY.register(target=[torch.ops.mlx.gguf_linear.default])
101+
def _gguf_linear_handler(P: MLXProgramBuilder, n: Node) -> Slot:
102+
"""Route ``mlx::gguf_linear`` lowering to the per-format implementation."""
103+
args = P.args(n)
104+
fmt = args[2] if len(args) >= 3 else None
105+
if fmt == "q6k":
106+
from executorch.backends.mlx.custom_kernel_ops.gguf.q6k.linear import (
107+
emit_linear,
108+
)
109+
110+
return emit_linear(P, n)
111+
raise NotImplementedError(
112+
f"mlx::gguf_linear: unsupported format {fmt!r}; only 'q6k' is supported"
113+
)

‎backends/mlx/custom_kernel_ops/gguf/q6k/__init__.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
#
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
#
8+
9+
"""GGUF **Q6_K** format implementation.
10+
11+
* :mod:`.common` -- shared primitives (constants, pure-torch dequant, Metal
12+
header). Re-exported here so ``from ...gguf.q6k import dequantize_q6_k`` stays
13+
lightweight (no MLX builder import).
14+
* :mod:`.linear` -- Q6_K mat-vec/mat-mat kernels + eager compute + lowering.
15+
* :mod:`.embedding` -- Q6_K gather kernel + eager compute + lowering.
16+
17+
The format-agnostic op routers live one level up in
18+
``custom_kernel_ops.gguf.linear`` / ``.embedding`` and dispatch here on
19+
``format``. ``.linear`` / ``.embedding`` are intentionally NOT imported here so
20+
importing :mod:`.common` for the pure-torch dequant does not pull in the builder.
21+
"""
22+
23+
from executorch.backends.mlx.custom_kernel_ops.gguf.q6k.common import ( # noqa: F401
24+
_Q6K_HEADER,
25+
dequantize_q6_k,
26+
Q6K_BLOCK_BYTES,
27+
QK_K,
28+
)

0 commit comments

Comments
 (0)