diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index 4822101bdca..086b45ff910 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -194,6 +194,52 @@ jobs: echo "Success: built $RUNNER" echo "::endgroup::" + # Build-only: the Muse Glimmer GGUF checkpoints are far too large to export in + # CI, so this covers the compile / link / metallib wiring of the runners -- + # which is what MLX backend changes are most likely to break. + test-mlx-muse-glimmer-build: + needs: run-decision + if: | + github.event_name == 'pull_request' || + needs.run-decision.outputs.is-full-run == 'true' + uses: pytorch/test-infra/.github/workflows/macos_job.yml@main + with: + default-packages: "" + job-name: test-mlx-muse-glimmer-build + runner: macos-14-xlarge + python-version: "3.12" + submodules: recursive + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 90 + script: | + set -eux + + echo "::group::Install ExecuTorch" + ${CONDA_RUN} python install_executorch.py > /dev/null + echo "::endgroup::" + + ${CONDA_RUN} pip list + + echo "::group::Build Muse Glimmer runners and worker with MLX" + ${CONDA_RUN} make muse-glimmer-mlx + BIN_DIR=cmake-out/examples/models/muse-glimmer + for bin in solo_runner dflash_runner muse_glimmer_worker; do + if [ ! -x "${BIN_DIR}/${bin}" ]; then + echo "Failed: ${bin} not found at ${BIN_DIR}/${bin}" + exit 1 + fi + done + # Runtime kernel load falls back to an absolute build-tree path, so a + # missing colocated metallib only breaks a relocated binary. Assert the + # copy so the runners stay self-contained (needs EXECUTORCH_BUILD_MLX=ON + # in the preset, which drives executorch_target_copy_mlx_metallib). + if [ ! -f "${BIN_DIR}/mlx.metallib" ]; then + echo "Failed: mlx.metallib not copied next to the runners" + exit 1 + fi + echo "Success: built solo_runner, dflash_runner and muse_glimmer_worker" + echo "::endgroup::" + backend-tester: needs: run-decision if: | diff --git a/Makefile b/Makefile index 444b4f08ee4..b6d69bc4024 100644 --- a/Makefile +++ b/Makefile @@ -91,7 +91,7 @@ # # ============================================================================== -.PHONY: voxtral-cuda voxtral-cpu voxtral-metal voxtral-mlx voxtral_realtime-cuda voxtral_realtime-cpu voxtral_realtime-metal voxtral_realtime-mlx voxtral_tts-cpu voxtral_tts-cuda whisper-cuda whisper-cuda-debug whisper-cpu whisper-metal parakeet-cuda parakeet-cuda-debug parakeet-cpu parakeet-metal parakeet-mlx parakeet-vulkan dinov2-cuda dinov2-cuda-debug sortformer-cuda sortformer-cpu silero-vad-cpu llama-cuda llama-cuda-debug llama-cpu lfm_2_5-mlx llava-cpu gemma3-cuda gemma3-cpu gemma4_31b-cuda gemma4_31b-mlx muse-glimmer-cuda qwen3_5_moe-cuda qwen3_5_moe-metal qwen3_5_moe-mlx clean help +.PHONY: voxtral-cuda voxtral-cpu voxtral-metal voxtral-mlx voxtral_realtime-cuda voxtral_realtime-cpu voxtral_realtime-metal voxtral_realtime-mlx voxtral_tts-cpu voxtral_tts-cuda whisper-cuda whisper-cuda-debug whisper-cpu whisper-metal parakeet-cuda parakeet-cuda-debug parakeet-cpu parakeet-metal parakeet-mlx parakeet-vulkan dinov2-cuda dinov2-cuda-debug sortformer-cuda sortformer-cpu silero-vad-cpu llama-cuda llama-cuda-debug llama-cpu lfm_2_5-mlx llava-cpu gemma3-cuda gemma3-cpu gemma4_31b-cuda gemma4_31b-mlx muse-glimmer-cuda muse-glimmer-mlx qwen3_5_moe-cuda qwen3_5_moe-metal qwen3_5_moe-mlx clean help help: @echo "This Makefile adds targets to build runners for various models on various backends. Run using \`make \`. Available targets:" @@ -130,6 +130,7 @@ help: @echo " gemma4_31b-cuda - Build Gemma 4 31B runner and worker with CUDA backend" @echo " gemma4_31b-mlx - Build Gemma 4 31B runner and worker with MLX backend" @echo " muse-glimmer-cuda - Build Muse Glimmer runners and worker with CUDA backend" + @echo " muse-glimmer-mlx - Build Muse Glimmer runners and worker with MLX backend" @echo " qwen3_5_moe-cuda - Build Qwen3.5 MoE runner with CUDA backend" @echo " qwen3_5_moe-metal - Build Qwen3.5 MoE runner with Metal backend" @echo " qwen3_5_moe-mlx - Build Qwen3.5 MoE runner with MLX backend" @@ -453,6 +454,17 @@ muse-glimmer-cuda: @echo " DFlash runner: cmake-out/examples/models/muse-glimmer/dflash_runner" @echo " Worker: cmake-out/examples/models/muse-glimmer/muse_glimmer_worker" +muse-glimmer-mlx: + @echo "==> Building and installing ExecuTorch with MLX..." + cmake --workflow --preset mlx-release + @echo "==> Building Muse Glimmer runners and worker with MLX..." + cd examples/models/muse-glimmer && cmake --workflow --preset muse-glimmer-mlx + @echo "" + @echo "✓ Build complete!" + @echo " Solo runner: cmake-out/examples/models/muse-glimmer/solo_runner" + @echo " DFlash runner: cmake-out/examples/models/muse-glimmer/dflash_runner" + @echo " Worker: cmake-out/examples/models/muse-glimmer/muse_glimmer_worker" + gemma4_31b-cuda: @echo "==> Building and installing ExecuTorch with CUDA..." cmake --workflow --preset llm-release-cuda diff --git a/backends/mlx/_memprofile.py b/backends/mlx/_memprofile.py new file mode 100644 index 00000000000..2a5dd2b5476 --- /dev/null +++ b/backends/mlx/_memprofile.py @@ -0,0 +1,188 @@ +# +# 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. +# + +"""Phase-attributed peak-memory profiling for MLX export. + +Enabled with ``ET_MLX_MEM_PROFILE=1``; inert otherwise:: + + ET_MLX_MEM_PROFILE=1 python -m ...export_dflash --target-gguf ... + + [mem] build: 42.84 GB -> 41.30 GB, -1.54 GB net, high-water 48.26 GB, \ +RAISED PEAK by 5.42 GB, 24.4s + +Reports macOS *physical footprint* -- what ``/usr/bin/time -l`` calls "peak +memory footprint" and Activity Monitor shows as "Memory". RSS is not a usable +proxy: a 16 GB model export peaked at 65 GB footprint versus 37 GB RSS. + +Only phases tagged ``RAISED PEAK`` can lower the number ``/usr/bin/time -l`` +reports; everything else is churn under the existing high-water mark. + +The watermark comes from the kernel rather than a sampling thread, which would +miss short-lived spikes -- a 2 GB allocation freed within 100 ms went unseen at +a 50 ms sampling interval. +""" + +from __future__ import annotations + +import ctypes +import ctypes.util +import os +import time +from contextlib import contextmanager +from typing import Optional + +from executorch.backends.mlx._logging import logger + +_ENV_VAR = "ET_MLX_MEM_PROFILE" + +# proc_pid_rusage(pid, flavor, rusage_info_t *buffer), from . +# Offsets into struct rusage_info_v4 (): a 16-byte ri_uuid +# followed by uint64 counters. Sanity-checked at runtime by _self_check(). +_RUSAGE_INFO_V4 = 4 +_OFF_PHYS_FOOTPRINT = 72 +_OFF_PROC_START_ABSTIME = 80 +_OFF_PROC_EXIT_ABSTIME = 88 +_OFF_LIFETIME_MAX_PHYS_FOOTPRINT = 240 +_RUSAGE_BUF_BYTES = 512 # generous; larger than any rusage_info_v4 + +# Bounds for the self-check: any live process is above the floor, and no real +# footprint approaches the ceiling. +_MIN_PLAUSIBLE = 1 << 20 # 1 MiB +_MAX_PLAUSIBLE = 1 << 50 # 1 PiB + +_libc = None +_usable: Optional[bool] = None +_depth = 0 + + +def _read_rusage() -> Optional[bytes]: + global _libc + try: + if _libc is None: + _libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True) + buf = (ctypes.c_uint8 * _RUSAGE_BUF_BYTES)() + if _libc.proc_pid_rusage(os.getpid(), _RUSAGE_INFO_V4, ctypes.byref(buf)) != 0: + return None + return bytes(buf) + except (AttributeError, OSError, ValueError): + return None + + +def _field(raw: Optional[bytes], offset: int) -> Optional[int]: + if raw is None: + return None + return int.from_bytes(raw[offset : offset + 8], "little") + + +def _self_check() -> bool: + """Smoke-test the hardcoded rusage offsets before trusting what they return. + + The struct layout is pinned by the flavor argument (Apple adds rusage_info_v5 + rather than reordering v4), so this is a guard against being wrong about the + layout, not against it changing underneath us. + + ri_proc_exit_abstime is the anchor: it is necessarily zero for the process + asking, so a non-zero read means the offsets are not landing where we think. + The bounds checks on the two footprint fields are weaker -- a misread can + still satisfy them by coincidence, since neighbouring counters hold numbers + of a similar magnitude -- so they catch gross errors only. + """ + raw = _read_rusage() + current = _field(raw, _OFF_PHYS_FOOTPRINT) + peak = _field(raw, _OFF_LIFETIME_MAX_PHYS_FOOTPRINT) + started = _field(raw, _OFF_PROC_START_ABSTIME) + exited = _field(raw, _OFF_PROC_EXIT_ABSTIME) + + if current is None: + reason = "proc_pid_rusage unavailable (non-macOS?)" + elif exited != 0 or not started: + reason = ( + f"struct layout unrecognized (proc_start={started}, proc_exit={exited}; " + "expected a non-zero start and a zero exit for a live process)" + ) + elif not _MIN_PLAUSIBLE <= current <= peak <= _MAX_PLAUSIBLE: + reason = f"implausible readings (current={current}, lifetime max={peak})" + else: + return True + + logger.warning(f"[mem] memory profiling disabled: {reason}") + return False + + +def enabled() -> bool: + """Whether profiling is switched on and the platform counters are trustworthy.""" + global _usable + if os.environ.get(_ENV_VAR, "0") == "0": + return False + if _usable is None: + _usable = _self_check() + return _usable + + +def phys_footprint() -> Optional[int]: + """Current physical footprint in bytes, or None if unavailable.""" + return _field(_read_rusage(), _OFF_PHYS_FOOTPRINT) + + +def peak_footprint() -> Optional[int]: + """Process lifetime maximum physical footprint in bytes. + + Matches the "peak memory footprint" line from ``/usr/bin/time -l``. + """ + return _field(_read_rusage(), _OFF_LIFETIME_MAX_PHYS_FOOTPRINT) + + +def _gb(n: Optional[int]) -> str: + return "?" if n is None else f"{n / (1 << 30):.2f} GB" + + +def _delta(before: Optional[int], after: Optional[int]) -> Optional[int]: + return None if (before is None or after is None) else after - before + + +@contextmanager +def mem_phase(name: str): + """Log footprint across `name`, attributing any new high-water mark to it. + + Nested phases are indented. Does nothing unless ET_MLX_MEM_PROFILE is set. + """ + global _depth + + if not enabled(): + yield + return + + indent = " " * _depth + start, start_peak = phys_footprint(), peak_footprint() + started_at = time.perf_counter() + _depth += 1 + try: + yield + finally: + _depth -= 1 + end, end_peak = phys_footprint(), peak_footprint() + net = _delta(start, end) + raised = _delta(start_peak, end_peak) + + sign = "+" if net is not None and net >= 0 else "" + message = ( + f"[mem]{indent} {name}: {_gb(start)} -> {_gb(end)}, " + f"{sign}{_gb(net)} net, high-water {_gb(end_peak)}" + ) + if raised: + message += f", RAISED PEAK by {_gb(raised)}" + logger.info(f"{message}, {time.perf_counter() - started_at:.1f}s") + + +def log_footprint(label: str) -> None: + """Log a one-off footprint reading.""" + if enabled(): + logger.info( + f"[mem] {label}: {_gb(phys_footprint())} " + f"(high-water {_gb(peak_footprint())})" + ) diff --git a/backends/mlx/builder/op_helpers.py b/backends/mlx/builder/op_helpers.py index c846513f708..48a2a0376ed 100644 --- a/backends/mlx/builder/op_helpers.py +++ b/backends/mlx/builder/op_helpers.py @@ -534,6 +534,28 @@ def emit_quantized_gather( ) +def mlx_qparams_supported( + in_features: int, num_groups: int, group_size: int, bits: int +) -> bool: + """Whether to_mlx_qparams + regroup_affine_scales can consume this layout. + + Mirrors their asserts using shape metadata only, so a handler can answer + "would I lower this?" without reading (and repacking) the weight: + + * to_mlx_qparams packs a row into whole uint32 words, so + ``in_features * bits`` must be a multiple of 32. + * regroup_affine_scales repeat-interleaves, which only ever splits a group + finer, so the weight's own group must be a whole multiple of the + MLX-legal ``group_size``. + """ + if in_features <= 0 or num_groups <= 0 or in_features % num_groups != 0: + return False + if (in_features * bits) % 32 != 0: + return False + weight_group_size = in_features // num_groups + return weight_group_size >= group_size and weight_group_size % group_size == 0 + + def to_mlx_qparams( qdata: torch.Tensor, scale: torch.Tensor, diff --git a/backends/mlx/builder/op_registry.py b/backends/mlx/builder/op_registry.py index 19668ca2c1b..5bdb6435d0d 100644 --- a/backends/mlx/builder/op_registry.py +++ b/backends/mlx/builder/op_registry.py @@ -8,7 +8,17 @@ from __future__ import annotations -from typing import Callable, Dict, List, Optional, Tuple, Type, TYPE_CHECKING, Union +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Tuple, + Type, + TYPE_CHECKING, + Union, +) from executorch.backends.mlx._logging import logger from torch.fx.node import Node @@ -23,6 +33,10 @@ ["MLXProgramBuilder", Node], Optional[Union["Slot", Tuple["Slot", ...]]] ] +# Support-check type: takes (builder, node) and returns whether the handler for +# that node will lower it. See PatternHandler.supported for the contract. +SupportCheck = Callable[["MLXProgramBuilder", Node], bool] + class PatternHandler: def __init__(self, head: Node, body: List[Node]) -> None: @@ -40,6 +54,32 @@ def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional[PatternHandle def __call__(self, P: MLXProgramBuilder, n: Node) -> None: raise NotImplementedError + def supported(self, P: MLXProgramBuilder, n: Node) -> bool: + """Optional: answer "would __call__ lower this node?" without running it. + + Support is normally decided by running the handler and seeing whether it + throws, but handlers that repack a quantized weight copy the whole tensor + to answer that question -- and the partitioner asks it several times per + export. Overriding this lets such a handler answer from node metadata + instead; it is then consulted in place of running the handler during + support checks (never during build(), which has to emit). + + Implementations must decide from ``n.meta`` alone and must never read + constant data, which is the cost being avoided. They must also agree with + __call__: saying yes to a node the handler then rejects makes + ops_to_not_decompose preserve an op that afterwards neither decomposes + nor lowers. Disagreement is caught by the op tests rather than a + dedicated check -- a false positive fails the export outright when + build() runs the handler for real, and a false negative shows up as a + missing delegate segment. + """ + raise NotImplementedError + + @classmethod + def has_support_check(cls) -> bool: + """Whether this class overrides supported().""" + return cls.supported is not PatternHandler.supported + def set_handlers(self, P: MLXProgramBuilder): if P.node_info[self.head].handler is not None: raise AssertionError( @@ -67,11 +107,13 @@ class MLXOpRegistry: def __init__(self): self._handlers: Dict[Union[str, Callable], Handler] = {} + self._support_checks: Dict[Union[str, Callable], SupportCheck] = {} self._patterns: Dict[str, Type[PatternHandler]] = {} def reset(self) -> None: """Reset the registry to empty state. Useful for testing.""" self._handlers.clear() + self._support_checks.clear() self._patterns.clear() def register(self, target: Union[str, Callable, list, tuple]): @@ -89,16 +131,44 @@ def deco(fn: Handler): def get_handler(self, node: Node) -> Optional[Handler]: """Get the handler for a node, or None if not registered.""" + return self._lookup(self._handlers, node) + + def register_support_check(self, target: Union[str, Callable, list, tuple]): + """Decorator registering a cheap support predicate for an op handler. + + The predicate takes (builder, node) and returns whether the handler will + lower the node. It is consulted instead of running the handler during + support checks, so it must decide from node metadata alone -- reading + constant data is exactly the cost being avoided. See + PatternHandler.supported for the full contract. + """ + + def deco(fn: SupportCheck): + targets = target if isinstance(target, (list, tuple)) else [target] + for t in targets: + if t in self._support_checks: + raise ValueError(f"Support check for {t} already registered") + self._support_checks[t] = fn + return fn + + return deco + + def get_support_check(self, node: Node) -> Optional[SupportCheck]: + """Get the support predicate for a node, or None if it has none.""" + return self._lookup(self._support_checks, node) + + @staticmethod + def _lookup(table: Dict[Union[str, Callable], Any], node: Node) -> Optional[Any]: t = node.target - if t in self._handlers: - return self._handlers[t] + if t in table: + return table[t] # Handle EdgeOpOverload by extracting the underlying ATen op - if hasattr(t, "_op") and t._op in self._handlers: - return self._handlers[t._op] + if hasattr(t, "_op") and t._op in table: + return table[t._op] # Check for string-based targets (e.g., higher_order ops) target_str = str(t) - if target_str in self._handlers: - return self._handlers[target_str] + if target_str in table: + return table[target_str] return None def registered_ops(self) -> set: @@ -118,6 +188,8 @@ def unregister(self, target: Union[str, Callable, list, tuple]) -> None: for t in targets: if t in self._handlers: del self._handlers[t] + if t in self._support_checks: + del self._support_checks[t] def register_pattern(self, name: str): """Decorator to register a pattern handler class.""" diff --git a/backends/mlx/builder/program_builder.py b/backends/mlx/builder/program_builder.py index 1cf94e9f1d1..829ebce121f 100644 --- a/backends/mlx/builder/program_builder.py +++ b/backends/mlx/builder/program_builder.py @@ -152,6 +152,13 @@ def __init__(self, ep: ExportedProgram, named_data_key_prefix: str = ""): # Unprefixed canonical-name → Slot for constants, populated by _build_io_maps(). # Used by get_named_data_store() to look up tensors without prefix interference. self._constant_name_to_slot: Dict[str, Slot] = {} + # Per-weight memo for repacking handlers, keyed by the raw weight's + # canonical name. Lets a weight feeding two ops (a tied embedding and + # lm_head, say) be repacked once, which also makes it safe to release + # the raw tensor after the first repack. + self.repack_cache: Dict[str, Any] = {} + # True while classifying nodes rather than emitting; gates weight release. + self._check_only: bool = False def _prefix_key(self, name: str) -> str: """Apply the named-data key prefix for the .pte namespace. @@ -290,6 +297,42 @@ def get_placeholder_target_and_tensor(self, node: Node) -> Tuple[str, torch.Tens raise KeyError(f"Unable to resolve placeholder {placeholder_name}") + def get_placeholder_target(self, node: Node) -> str: + """Resolve a placeholder to its state_dict / constants key. + + Unlike get_placeholder_target_and_tensor this never touches the data, so + it still answers after the tensor has been released. + """ + assert node.op == "placeholder" + for ispec in self.ep.graph_signature.input_specs: + if ispec.arg.name == node.name and ispec.target is not None: + return ispec.target + raise KeyError(f"Unable to resolve placeholder {node.name}") + + def release_placeholder_tensor(self, node: Node) -> None: + """Drop a placeholder's data once a handler has finished repacking it. + + A repacking handler replaces a raw weight with new constants, leaving + the original dead. It would otherwise stay in the ExportedProgram until + get_named_data_store() sweeps unused constants -- which only runs after + the whole graph is built, so every raw weight and every repacked weight + are live simultaneously (+16 GB on a 16 GB Q4_K model). + + Only call this when the raw data is fully consumed: the lowering must + not reference the placeholder's slot in any emitted instruction, or it + will have no data to serialize (get_named_data_store raises in that + case rather than emitting a .pte with a missing weight). + + Ignored during a support check. Those run against the caller's own + program rather than a delegate submodule, and torch still needs its + state dict afterwards -- deleting there fails later in + _unlift_exported_program_lifted_states. The check-only builder is + discarded anyway, so there is nothing to reclaim. + """ + if self._check_only: + return + self._delete_constant_tensor(self.get_placeholder_target(node)) + def slot_to_tid(self, slot: Slot) -> Tid: """Convert a tensor Slot to a Tid, recording it for later remapping.""" assert slot.id_type == IdType.Tensor @@ -523,7 +566,7 @@ def _apply_patterns(self) -> None: for handler in matcher.find_patterns(): handler.set_handlers(self) - def _process_nodes(self) -> None: # noqa C901 + def _process_nodes(self, check_only: bool = False) -> None: # noqa C901 """ Common logic for processing all nodes: create slots, match patterns, run handlers. @@ -536,7 +579,12 @@ def _process_nodes(self) -> None: # noqa C901 The ordering is important: patterns must be matched before noops because some pattern body nodes (e.g., update_cache) have no users since they mutate in-place, but they're not dead - they're handled by the pattern. + + With ``check_only``, a handler that provides a support check is asked + instead of being run. Only support status is populated in that mode; no + instructions are emitted for those nodes and no constants are repacked. """ + self._check_only = check_only self._make_io_slots() # Apply patterns BEFORE _mark_noop so pattern body nodes don't get @@ -545,12 +593,31 @@ def _process_nodes(self) -> None: # noqa C901 self._apply_patterns() self._mark_noop() + def mark_unsupported(n: Node, reason: str) -> None: + if n.meta.get("val", None) is not None: + self.slot_manager.make_or_get_slots(n) + self._mark_unsupported(n, reason) + for n in self.ep.graph.nodes: if self._is_handled(n): continue if self.node_info[n].handler is not None: handler = self.node_info[n].handler + if ( + check_only + and isinstance(handler, PatternHandler) + and type(handler).has_support_check() + ): + if handler.supported(self, n): + self._mark_supported(n, handler=handler) + else: + mark_unsupported( + n, + f"{type(handler).__name__}.supported() rejected " + f"target={n.target}", + ) + continue with self.tmp_scope(): handler(self, n) self._mark_supported(n, handler=handler) @@ -559,9 +626,7 @@ def _process_nodes(self) -> None: # noqa C901 # Check input dtypes before processing node unsupported_dtype_msg = _check_input_dtypes(n) if unsupported_dtype_msg is not None: - if n.meta.get("val", None) is not None: - self.slot_manager.make_or_get_slots(n) - self._mark_unsupported(n, unsupported_dtype_msg) + mark_unsupported(n, unsupported_dtype_msg) continue if n.op in ("placeholder", "output"): @@ -591,12 +656,20 @@ def _process_nodes(self) -> None: # noqa C901 handler = REGISTRY.get_handler(n) if handler is None: - msg = f"no handler for target={n.target}" - if n.meta.get("val", None) is not None: - self.slot_manager.make_or_get_slots(n) - self._mark_unsupported(n, msg) + mark_unsupported(n, f"no handler for target={n.target}") continue + if check_only: + check = REGISTRY.get_support_check(n) + if check is not None: + if check(self, n): + self._mark_supported(n, handler=handler) + else: + mark_unsupported( + n, f"{check.__name__} rejected target={n.target}" + ) + continue + try: with self.tmp_scope(): handler(self, n) @@ -604,9 +677,7 @@ def _process_nodes(self) -> None: # noqa C901 except Exception as e: trace_str = traceback.format_exc() msg = f"{handler} failed for {n.target}: {e}.\n{trace_str}" - if n.meta.get("val", None) is not None: - self.slot_manager.make_or_get_slots(n) - self._mark_unsupported(n, msg) + mark_unsupported(n, msg) def check_support_only(self) -> None: """ @@ -618,8 +689,12 @@ def check_support_only(self) -> None: Use this method for ops_to_not_decompose() and similar queries where you only need to know support status, not the full compiled graph. + + Handlers that provide a support check are asked rather than run, so they + do not repack their weights here; the rest are still run, because + whether a handler throws is the only way to know if it can lower a node. """ - self._process_nodes() + self._process_nodes(check_only=True) # NOTE: We intentionally skip _verify_build() and _build_mlx_graph() here # because _build_mlx_graph() calls int() on tensor shapes which evaluates # SymInts and corrupts the shape_env. This method is used for @@ -1075,7 +1150,15 @@ def get_named_data_store(self) -> NamedDataStore: for canonical_name, _slot in entries: tensor = self._find_constant_tensor(canonical_name) if tensor is None: - continue + # Every entry here is a constant the graph references, so missing + # data means a handler released a weight it still needed (see + # release_placeholder_tensor). Serializing anyway would emit a + # .pte with a silently missing weight. + raise RuntimeError( + f"No data for constant '{canonical_name}', which the MLX " + f"graph references. A handler likely released it while a " + f"lowering still needed it." + ) t = tensor.detach().cpu().contiguous() named_data_store.add_named_data( diff --git a/backends/mlx/custom_kernel_ops/gguf/patterns.py b/backends/mlx/custom_kernel_ops/gguf/patterns.py index 0632ceb27f9..7147a767ac0 100644 --- a/backends/mlx/custom_kernel_ops/gguf/patterns.py +++ b/backends/mlx/custom_kernel_ops/gguf/patterns.py @@ -40,6 +40,9 @@ from executorch.backends.mlx.builder.op_registry import PatternHandler, REGISTRY from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder from executorch.backends.mlx.builder.slot_manager import Slot +from executorch.backends.mlx.custom_kernel_ops.gguf.q4k.common import Q4K_BLOCK_BYTES +from executorch.backends.mlx.custom_kernel_ops.gguf.q5k.common import Q5K_BLOCK_BYTES +from executorch.backends.mlx.custom_kernel_ops.gguf.q6k.common import Q6K_BLOCK_BYTES from executorch.backends.mlx.pattern_utils import has_single_user, match_target from torch.export.exported_program import ExportedProgram from torch.fx.node import Node @@ -49,6 +52,29 @@ _LINEAR_TYPES = {"q4_k", "q5_k", "q6_k"} _EMBEDDING_TYPES = {"q4_k", "q5_k", "q6_k"} +_BLOCK_BYTES = { + "q4_k": Q4K_BLOCK_BYTES, + "q5_k": Q5K_BLOCK_BYTES, + "q6_k": Q6K_BLOCK_BYTES, +} + + +def _blob_lowers(weight_node: Node, ggml_type: str) -> bool: + """Whether the raw GGUF blob has a shape every lowering path can consume. + + Mirrors the checks each ``emit_*`` makes before touching the blob: 2-D + ``(rows, row_bytes)``, statically shaped, a whole number of super-blocks per + row. Deliberately shape-only, so it holds for both the fused-kernel and the + MLX-native repack path (which one runs depends on ET_MLX_EMIT_DIRECT_GGUF). + """ + val = weight_node.meta.get("val", None) + if val is None or val.dim() != 2: + return False + rows, row_bytes = val.shape + if not isinstance(rows, int) or not isinstance(row_bytes, int): + return False + return row_bytes % _BLOCK_BYTES[ggml_type] == 0 + def parse_dequantize_gguf_node( node: Node, @@ -108,6 +134,9 @@ def maybe_create(cls, ep: ExportedProgram, head: Node): return None return cls(head, [dequant], weight, ggml_type, output_dtype) + def supported(self, P: MLXProgramBuilder, n: Node) -> bool: + return _blob_lowers(self.weight, self.ggml_type) + def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: assert n == self.head x_node = n.args[0] @@ -159,6 +188,9 @@ def maybe_create(cls, ep: ExportedProgram, head: Node): return None return cls(head, [dequant], weight, ggml_type, output_dtype) + def supported(self, P: MLXProgramBuilder, n: Node) -> bool: + return _blob_lowers(self.weight, self.ggml_type) + def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: assert n == self.head indices_node = n.args[1] diff --git a/backends/mlx/custom_kernel_ops/gguf/q4k/embedding_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q4k/embedding_mlx_native.py index 3a032c2b7fd..2388dc18aa6 100644 --- a/backends/mlx/custom_kernel_ops/gguf/q4k/embedding_mlx_native.py +++ b/backends/mlx/custom_kernel_ops/gguf/q4k/embedding_mlx_native.py @@ -19,10 +19,7 @@ from executorch.backends.mlx.builder.op_helpers import emit_quantized_gather from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder from executorch.backends.mlx.builder.slot_manager import Slot -from executorch.backends.mlx.custom_kernel_ops.gguf.q4k.repack_mlx import ( - _BITS, - repack_mlx, -) +from executorch.backends.mlx.custom_kernel_ops.gguf.repack_mlx import Q4_K, repack_mlx from torch.fx.node import Node @@ -40,7 +37,7 @@ def emit_embedding( embedding. """ w_slot, scales_slot, biases_slot, group_size = repack_mlx( - P, weight_node, scale_dtype=output_dtype + P, weight_node, Q4_K, scale_dtype=output_dtype ) (indices_slot,) = P.slot_map([indices_node]) @@ -53,7 +50,7 @@ def emit_embedding( scales_slot, biases_slot, group_size=group_size, - bits=_BITS, + bits=Q4_K.bits, mode="affine", out_dtype=output_dtype, ) diff --git a/backends/mlx/custom_kernel_ops/gguf/q4k/linear_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q4k/linear_mlx_native.py index a746948b4c0..3317047a285 100644 --- a/backends/mlx/custom_kernel_ops/gguf/q4k/linear_mlx_native.py +++ b/backends/mlx/custom_kernel_ops/gguf/q4k/linear_mlx_native.py @@ -21,10 +21,7 @@ from executorch.backends.mlx.builder.op_helpers import torch_dtype_to_scalar_type from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder from executorch.backends.mlx.builder.slot_manager import Slot -from executorch.backends.mlx.custom_kernel_ops.gguf.q4k.repack_mlx import ( - _BITS, - repack_mlx, -) +from executorch.backends.mlx.custom_kernel_ops.gguf.repack_mlx import Q4_K, repack_mlx from executorch.backends.mlx.serialization.mlx_graph_schema import ( AddNode, AsTypeNode, @@ -47,7 +44,7 @@ def emit_linear( MLX-format constants are serialized. """ w_slot, scales_slot, biases_slot, group_size = repack_mlx( - P, weight_node, scale_dtype=x_node.meta["val"].dtype + P, weight_node, Q4_K, scale_dtype=x_node.meta["val"].dtype ) x_slot, bias_slot = P.slot_map([x_node, bias_node]) @@ -60,7 +57,7 @@ def emit_linear( biases=P.slot_to_tid(biases_slot), out=P.slot_to_tid(out), group_size=group_size, - bits=_BITS, + bits=Q4_K.bits, mode="affine", transpose=True, ) diff --git a/backends/mlx/custom_kernel_ops/gguf/q4k/repack_mlx.py b/backends/mlx/custom_kernel_ops/gguf/q4k/repack_mlx.py deleted file mode 100644 index 16c8543d098..00000000000 --- a/backends/mlx/custom_kernel_ops/gguf/q4k/repack_mlx.py +++ /dev/null @@ -1,60 +0,0 @@ -# -# 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. -# - -"""Q4_K -> MLX qparam repack for the MLX-native lowering path. - -Used when ``ET_MLX_EMIT_DIRECT_GGUF=0``: the GGUF blob is unpacked and repacked -into MLX affine 4-bit qparams at export time instead of being consumed directly -by fused Metal kernels. -""" - -from __future__ import annotations - -from typing import Optional, Tuple - -import torch - -from executorch.backends.mlx.builder.op_helpers import to_mlx_qparams -from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder -from executorch.backends.mlx.builder.slot_manager import Slot -from torch.fx.node import Node - -_BITS = 4 - - -def repack_mlx( - P: MLXProgramBuilder, - weight_node: Node, - scale_dtype: Optional[torch.dtype] = None, -) -> Tuple[Slot, Slot, Slot, int]: - """Unpack a raw Q4_K blob and repack into MLX qparam constants. - - Adjacent sub-blocks with identical scale/min are merged into a larger group - size (up to 128) when lossless, so ``group_size`` may be 32, 64, or 128. - Returns ``(packed_slot, scales_slot, biases_slot, group_size)``. - - ``scale_dtype`` sets the dtype of the emitted scales/biases constants; pass - the activation dtype so MLX ``quantized_matmul`` does not promote (a bf16 - activation with f16 scales, or vice versa, promotes to float32). - """ - from executorch.extension.llm.export.gguf import ExportableGGUFTensor - - weight_target, raw = P.get_placeholder_target_and_tensor(weight_node) - intx = ExportableGGUFTensor.from_raw(raw, "q4_k").to_intx_unpacked_to_int8_tensor( - max_group_size=128, scale_dtype=scale_dtype - ) - group_size = int(intx.block_size[-1]) - qdata, scale, zero_point = intx.qdata, intx.scale, intx.zero_point - del intx # drop the tensor-subclass wrapper; keep only the fields we need - packed, biases = to_mlx_qparams(qdata, scale, zero_point, _BITS) - del qdata # the (N, K) int8 is no longer needed once packed - - packed_slot = P.make_or_get_constant(f"{weight_target}_q4k_packed", packed) - scales_slot = P.make_or_get_constant(f"{weight_target}_q4k_scales", scale) - biases_slot = P.make_or_get_constant(f"{weight_target}_q4k_biases", biases) - return packed_slot, scales_slot, biases_slot, group_size diff --git a/backends/mlx/custom_kernel_ops/gguf/q5k/embedding_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q5k/embedding_mlx_native.py index 52f796afd16..9b3f33edf57 100644 --- a/backends/mlx/custom_kernel_ops/gguf/q5k/embedding_mlx_native.py +++ b/backends/mlx/custom_kernel_ops/gguf/q5k/embedding_mlx_native.py @@ -19,10 +19,7 @@ from executorch.backends.mlx.builder.op_helpers import emit_quantized_gather from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder from executorch.backends.mlx.builder.slot_manager import Slot -from executorch.backends.mlx.custom_kernel_ops.gguf.q5k.repack_mlx import ( - _BITS, - repack_mlx, -) +from executorch.backends.mlx.custom_kernel_ops.gguf.repack_mlx import Q5_K, repack_mlx from torch.fx.node import Node @@ -40,7 +37,7 @@ def emit_embedding( embedding. """ w_slot, scales_slot, biases_slot, group_size = repack_mlx( - P, weight_node, scale_dtype=output_dtype + P, weight_node, Q5_K, scale_dtype=output_dtype ) (indices_slot,) = P.slot_map([indices_node]) @@ -53,7 +50,7 @@ def emit_embedding( scales_slot, biases_slot, group_size=group_size, - bits=_BITS, + bits=Q5_K.bits, mode="affine", out_dtype=output_dtype, ) diff --git a/backends/mlx/custom_kernel_ops/gguf/q5k/linear_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q5k/linear_mlx_native.py index e5543fbc7f5..2bb0744e117 100644 --- a/backends/mlx/custom_kernel_ops/gguf/q5k/linear_mlx_native.py +++ b/backends/mlx/custom_kernel_ops/gguf/q5k/linear_mlx_native.py @@ -21,10 +21,7 @@ from executorch.backends.mlx.builder.op_helpers import torch_dtype_to_scalar_type from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder from executorch.backends.mlx.builder.slot_manager import Slot -from executorch.backends.mlx.custom_kernel_ops.gguf.q5k.repack_mlx import ( - _BITS, - repack_mlx, -) +from executorch.backends.mlx.custom_kernel_ops.gguf.repack_mlx import Q5_K, repack_mlx from executorch.backends.mlx.serialization.mlx_graph_schema import ( AddNode, AsTypeNode, @@ -47,7 +44,7 @@ def emit_linear( MLX-format constants are serialized. """ w_slot, scales_slot, biases_slot, group_size = repack_mlx( - P, weight_node, scale_dtype=x_node.meta["val"].dtype + P, weight_node, Q5_K, scale_dtype=x_node.meta["val"].dtype ) x_slot, bias_slot = P.slot_map([x_node, bias_node]) @@ -60,7 +57,7 @@ def emit_linear( biases=P.slot_to_tid(biases_slot), out=P.slot_to_tid(out), group_size=group_size, - bits=_BITS, + bits=Q5_K.bits, mode="affine", transpose=True, ) diff --git a/backends/mlx/custom_kernel_ops/gguf/q5k/repack_mlx.py b/backends/mlx/custom_kernel_ops/gguf/q5k/repack_mlx.py deleted file mode 100644 index 7477fc68da0..00000000000 --- a/backends/mlx/custom_kernel_ops/gguf/q5k/repack_mlx.py +++ /dev/null @@ -1,60 +0,0 @@ -# -# 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. -# - -"""Q5_K -> MLX qparam repack for the MLX-native lowering path. - -Used when ``ET_MLX_EMIT_DIRECT_GGUF=0``: the GGUF blob is unpacked and repacked -into MLX affine 5-bit qparams at export time instead of being consumed directly -by fused Metal kernels. -""" - -from __future__ import annotations - -from typing import Optional, Tuple - -import torch - -from executorch.backends.mlx.builder.op_helpers import to_mlx_qparams -from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder -from executorch.backends.mlx.builder.slot_manager import Slot -from torch.fx.node import Node - -_BITS = 5 - - -def repack_mlx( - P: MLXProgramBuilder, - weight_node: Node, - scale_dtype: Optional[torch.dtype] = None, -) -> Tuple[Slot, Slot, Slot, int]: - """Unpack a raw Q5_K blob and repack into MLX qparam constants. - - Adjacent sub-blocks with identical scale/min are merged into a larger group - size (up to 128) when lossless, so ``group_size`` may be 32, 64, or 128. - Returns ``(packed_slot, scales_slot, biases_slot, group_size)``. - - ``scale_dtype`` sets the dtype of the emitted scales/biases constants; pass - the activation dtype so MLX ``quantized_matmul`` does not promote (a bf16 - activation with f16 scales, or vice versa, promotes to float32). - """ - from executorch.extension.llm.export.gguf import ExportableGGUFTensor - - weight_target, raw = P.get_placeholder_target_and_tensor(weight_node) - intx = ExportableGGUFTensor.from_raw(raw, "q5_k").to_intx_unpacked_to_int8_tensor( - max_group_size=128, scale_dtype=scale_dtype - ) - group_size = int(intx.block_size[-1]) - qdata, scale, zero_point = intx.qdata, intx.scale, intx.zero_point - del intx # drop the tensor-subclass wrapper; keep only the fields we need - packed, biases = to_mlx_qparams(qdata, scale, zero_point, _BITS) - del qdata # the (N, K) int8 is no longer needed once packed - - packed_slot = P.make_or_get_constant(f"{weight_target}_q5k_packed", packed) - scales_slot = P.make_or_get_constant(f"{weight_target}_q5k_scales", scale) - biases_slot = P.make_or_get_constant(f"{weight_target}_q5k_biases", biases) - return packed_slot, scales_slot, biases_slot, group_size diff --git a/backends/mlx/custom_kernel_ops/gguf/q6k/embedding_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q6k/embedding_mlx_native.py index 4ccf56c422a..d3c4a6d7c91 100644 --- a/backends/mlx/custom_kernel_ops/gguf/q6k/embedding_mlx_native.py +++ b/backends/mlx/custom_kernel_ops/gguf/q6k/embedding_mlx_native.py @@ -23,10 +23,7 @@ from executorch.backends.mlx.builder.op_helpers import emit_quantized_gather from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder from executorch.backends.mlx.builder.slot_manager import Slot -from executorch.backends.mlx.custom_kernel_ops.gguf.q6k.repack_mlx import ( - _BITS, - repack_mlx, -) +from executorch.backends.mlx.custom_kernel_ops.gguf.repack_mlx import Q6_K, repack_mlx from torch.fx.node import Node @@ -42,7 +39,7 @@ def emit_embedding( Returns the output slot, or ``None`` when the weight does not merge to an MLX-supported group size (the caller should fall back to the fused gather). """ - repacked = repack_mlx(P, weight_node, scale_dtype=output_dtype) + repacked = repack_mlx(P, weight_node, Q6_K, scale_dtype=output_dtype) if repacked is None: return None w_slot, scales_slot, biases_slot, group_size = repacked @@ -57,7 +54,7 @@ def emit_embedding( scales_slot, biases_slot, group_size=group_size, - bits=_BITS, + bits=Q6_K.bits, mode="affine", out_dtype=output_dtype, ) diff --git a/backends/mlx/custom_kernel_ops/gguf/q6k/linear_mlx_native.py b/backends/mlx/custom_kernel_ops/gguf/q6k/linear_mlx_native.py index 1672754e690..771f9010e35 100644 --- a/backends/mlx/custom_kernel_ops/gguf/q6k/linear_mlx_native.py +++ b/backends/mlx/custom_kernel_ops/gguf/q6k/linear_mlx_native.py @@ -22,10 +22,7 @@ from executorch.backends.mlx.builder.op_helpers import torch_dtype_to_scalar_type from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder from executorch.backends.mlx.builder.slot_manager import Slot -from executorch.backends.mlx.custom_kernel_ops.gguf.q6k.repack_mlx import ( - _BITS, - repack_mlx, -) +from executorch.backends.mlx.custom_kernel_ops.gguf.repack_mlx import Q6_K, repack_mlx from executorch.backends.mlx.serialization.mlx_graph_schema import ( AddNode, AsTypeNode, @@ -46,7 +43,7 @@ def emit_linear( Returns the output slot, or ``None`` when the weight does not merge to an MLX-supported group size (the caller should fall back to fused kernels). """ - repacked = repack_mlx(P, weight_node, scale_dtype=x_node.meta["val"].dtype) + repacked = repack_mlx(P, weight_node, Q6_K, scale_dtype=x_node.meta["val"].dtype) if repacked is None: return None w_slot, scales_slot, biases_slot, group_size = repacked @@ -61,7 +58,7 @@ def emit_linear( biases=P.slot_to_tid(biases_slot), out=P.slot_to_tid(out), group_size=group_size, - bits=_BITS, + bits=Q6_K.bits, mode="affine", transpose=True, ) diff --git a/backends/mlx/custom_kernel_ops/gguf/q6k/repack_mlx.py b/backends/mlx/custom_kernel_ops/gguf/q6k/repack_mlx.py deleted file mode 100644 index c2372e338a2..00000000000 --- a/backends/mlx/custom_kernel_ops/gguf/q6k/repack_mlx.py +++ /dev/null @@ -1,87 +0,0 @@ -# -# 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. -# - -"""Q6_K -> MLX qparam repack for the MLX-native lowering path. - -Used when ``ET_MLX_EMIT_DIRECT_GGUF=0``: the GGUF blob is unpacked and repacked -into MLX affine 6-bit qparams at export time instead of being consumed directly -by fused Metal kernels. - -Q6_K's native group size is 16, which MLX's affine kernels do not support (only -32/64/128), so the repack succeeds only when adjacent sub-blocks merge -losslessly into a group size >= 32 (see -:func:`...gguf.ExportableGGUFTensor.to_intx_unpacked_to_int8_tensor`). When it -cannot, :func:`repack_mlx` returns ``None`` so the caller falls back to the -fused kernels. -""" - -from __future__ import annotations - -from typing import Optional, Tuple - -import torch - -from executorch.backends.mlx.builder.op_helpers import ( - emit_quantized_biases, - to_mlx_qparams, -) -from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder -from executorch.backends.mlx.builder.slot_manager import Slot -from torch.fx.node import Node - -_BITS = 6 - -# MLX affine kernels only support these group sizes; Q6_K's native 16 is not one. -_MIN_MLX_GROUP_SIZE = 32 - - -def repack_mlx( - P: MLXProgramBuilder, - weight_node: Node, - scale_dtype: Optional[torch.dtype] = None, -) -> Optional[Tuple[Slot, Slot, Slot, int]]: - """Unpack a raw Q6_K blob and repack into MLX qparam constants. - - Adjacent sub-blocks with identical scale are merged into a larger group size - (up to 128) when lossless. Returns ``(packed_slot, scales_slot, biases_slot, - group_size)`` when the merged ``group_size`` is MLX-compatible (>= 32), or - ``None`` when it is not (so the caller can fall back to fused kernels). - - ``scale_dtype`` sets the dtype of the emitted scales/biases constants; pass - the activation dtype so MLX ``quantized_matmul`` does not promote (a bf16 - activation with f16 scales, or vice versa, promotes to float32). - """ - from executorch.extension.llm.export.gguf import ExportableGGUFTensor - - weight_target, raw = P.get_placeholder_target_and_tensor(weight_node) - intx = ExportableGGUFTensor.from_raw(raw, "q6_k").to_intx_unpacked_to_int8_tensor( - max_group_size=128, scale_dtype=scale_dtype - ) - group_size = int(intx.block_size[-1]) - if group_size < _MIN_MLX_GROUP_SIZE: - return None - - qdata, scale, zero_point = intx.qdata, intx.scale, intx.zero_point - del intx # drop the tensor-subclass wrapper; keep only the fields we need - packed, biases = to_mlx_qparams(qdata, scale, zero_point, _BITS) - del qdata # the (N, K) int8 is no longer needed once packed - - packed_slot = P.make_or_get_constant(f"{weight_target}_q6k_packed", packed) - scales_slot = P.make_or_get_constant(f"{weight_target}_q6k_scales", scale) - # Q6_K is symmetric (zero-point 0): emit_quantized_biases computes - # biases = -scale * 2^(bits-1) in the init chain instead of serializing them. - biases_slot = emit_quantized_biases( - P, - f"{weight_target}_q6k", - scale, - zero_point, - _BITS, - biases, - scales_slot, - ) - return packed_slot, scales_slot, biases_slot, group_size diff --git a/backends/mlx/custom_kernel_ops/gguf/repack_mlx.py b/backends/mlx/custom_kernel_ops/gguf/repack_mlx.py new file mode 100644 index 00000000000..d5442a31ca9 --- /dev/null +++ b/backends/mlx/custom_kernel_ops/gguf/repack_mlx.py @@ -0,0 +1,132 @@ +# +# 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. +# + +"""Shared GGUF -> MLX qparam repack for the MLX-native lowering path. + +Used when ``ET_MLX_EMIT_DIRECT_GGUF=0``: the GGUF blob is unpacked and repacked +into MLX affine qparams at export time instead of being consumed directly by +fused Metal kernels. + +Q4_K / Q5_K / Q6_K repack identically -- they differ only in bit width and in +whether the merged group size can land below what MLX supports -- so each is +described by a :class:`GgufMlxFormat` and shares one implementation. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch + +from executorch.backends.mlx.builder.op_helpers import ( + emit_quantized_biases, + to_mlx_qparams, +) +from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder +from executorch.backends.mlx.builder.slot_manager import Slot +from torch.fx.node import Node + +# MLX's affine kernels only implement group sizes 32 / 64 / 128. +_MLX_MIN_GROUP_SIZE = 32 + +# Adjacent sub-blocks with an identical scale (and min) are merged into a larger +# group when lossless, up to this size. +_MAX_GROUP_SIZE = 128 + + +@dataclass(frozen=True) +class GgufMlxFormat: + """Per-format parameters for the MLX-native repack path.""" + + ggml_type: str + bits: int + # Q6_K's native sub-blocks are 16 wide -- below _MLX_MIN_GROUP_SIZE -- so a + # weight whose sub-blocks do not merge losslessly cannot use this path at + # all and the caller has to fall back to the fused kernels. Q4_K / Q5_K are + # natively 32, so they always clear the minimum. + may_fall_below_min_group: bool = False + + @property + def key_prefix(self) -> str: + """Prefix for the constants this format emits, e.g. ``q4_k`` -> ``q4k``.""" + return self.ggml_type.replace("_", "") + + +Q4_K = GgufMlxFormat("q4_k", 4) +Q5_K = GgufMlxFormat("q5_k", 5) +Q6_K = GgufMlxFormat("q6_k", 6, may_fall_below_min_group=True) + + +def repack_mlx( + P: MLXProgramBuilder, + weight_node: Node, + fmt: GgufMlxFormat, + scale_dtype: Optional[torch.dtype] = None, +) -> Optional[Tuple[Slot, Slot, Slot, int]]: + """Unpack a raw GGUF blob and repack it into MLX qparam constants. + + Returns ``(packed_slot, scales_slot, biases_slot, group_size)``, where + ``group_size`` is the merged size (32, 64 or 128). Returns ``None`` only for + a format with ``may_fall_below_min_group`` whose weight does not merge to an + MLX-supported group size, so the caller can fall back to fused kernels. + + ``scale_dtype`` sets the dtype of the emitted scales/biases constants; pass + the activation dtype so MLX ``quantized_matmul`` does not promote (a bf16 + activation with f16 scales, or vice versa, promotes to float32). + + Once this path is committed to, the raw blob is released so it does not sit + alongside the packed constants for the rest of the build. It is deliberately + kept when returning ``None``: the fused-kernel fallback reads those bytes. + """ + from executorch.extension.llm.export.gguf import ExportableGGUFTensor + + weight_target = P.get_placeholder_target(weight_node) + cached = P.repack_cache.get(weight_target) + if cached is not None: + # Second consumer of a shared weight (e.g. a tied embedding / lm_head). + # The raw blob is already gone, so it must not be read again. + return cached + + _, raw = P.get_placeholder_target_and_tensor(weight_node) + intx = ExportableGGUFTensor.from_raw( + raw, fmt.ggml_type + ).to_intx_unpacked_to_int8_tensor( + max_group_size=_MAX_GROUP_SIZE, scale_dtype=scale_dtype + ) + del raw + + group_size = int(intx.block_size[-1]) + if group_size < _MLX_MIN_GROUP_SIZE: + if fmt.may_fall_below_min_group: + # Falling back to the fused kernels, which consume the raw blob. + return None + raise ValueError( + f"{fmt.ggml_type}: merged group size {group_size} is below MLX's " + f"minimum of {_MLX_MIN_GROUP_SIZE}" + ) + + P.release_placeholder_tensor(weight_node) + qdata, scale, zero_point = intx.qdata, intx.scale, intx.zero_point + del intx # drop the tensor-subclass wrapper; keep only the fields we need + packed, biases = to_mlx_qparams(qdata, scale, zero_point, fmt.bits) + del qdata # the (N, K) int8 is no longer needed once packed + + key = f"{weight_target}_{fmt.key_prefix}" + packed_slot = P.make_or_get_constant(f"{key}_packed", packed) + scales_slot = P.make_or_get_constant(f"{key}_scales", scale) + # A symmetric format (zero-point 0, e.g. Q6_K) needs no serialized biases: + # emit_quantized_biases computes them as -scale * 2^(bits-1) in the init + # chain instead, and falls back to a constant when the zero-point is not 0. + biases_slot = emit_quantized_biases( + P, key, scale, zero_point, fmt.bits, biases, scales_slot + ) + + result = (packed_slot, scales_slot, biases_slot, group_size) + P.repack_cache[weight_target] = result + return result diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py index ff2b697c216..05fa1e3bedc 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -1863,6 +1863,31 @@ def _gather_mm_handler(P: MLXProgramBuilder, n: Node) -> Slot: return out +@REGISTRY.register_support_check(target=[torch.ops.mlx.gather_qmm.default]) +def _gather_qmm_supported(P: MLXProgramBuilder, n: Node) -> bool: + """Whether _gather_qmm_handler can repack these expert weights. + + Mirrors the to_mlx_qparams asserts against the [E, out, in] weight's + metadata rather than the weight itself -- MoE experts are the largest + constants in the model, so repacking them just to answer a support query is + the most expensive way to ask. + """ + w_node = n.args[1] if len(n.args) > 1 else None + if not isinstance(w_node, Node): + return False + w = w_node.meta.get("val", None) + if w is None or w.dim() != 3: + return False + cols = w.shape[-1] + if not isinstance(cols, int): + return False + bits = n.args[8] if len(n.args) > 8 else n.kwargs.get("bits", 4) + if w.dtype == torch.uint8: + # Prepacked nibbles are viewed straight to uint32, two values per byte. + return bits == 4 and cols % 4 == 0 + return w.dtype == torch.int8 and (cols * bits) % 32 == 0 + + @REGISTRY.register(target=[torch.ops.mlx.gather_qmm.default]) def _gather_qmm_handler(P: MLXProgramBuilder, n: Node) -> Slot: """Handle mlx::gather_qmm — fused gather + dequant + matmul for quantized MoE experts. diff --git a/backends/mlx/partitioner.py b/backends/mlx/partitioner.py index 0896cafc301..7814e883588 100644 --- a/backends/mlx/partitioner.py +++ b/backends/mlx/partitioner.py @@ -16,10 +16,12 @@ from __future__ import annotations import inspect +import weakref from typing import Any, Callable, Dict, List, Tuple, Union import torch from executorch.backends.mlx._logging import logger +from executorch.backends.mlx._memprofile import mem_phase from executorch.backends.mlx.preprocess import MLXBackend from executorch.exir.backend.backend_details import CompileSpec from executorch.exir.backend.canonical_partitioners.pattern_op_partitioner import ( @@ -43,6 +45,11 @@ class MLXOperatorSupport(OperatorSupportBase): Uses MLXProgramBuilder to determine support - this ensures the partitioner uses the exact same logic as the actual compilation. A node is supported if the builder can handle it (either via direct handler or pattern match). + + The builder's verdicts are copied out and the builder itself dropped: + running the handlers repacks every quantized weight into builder-owned + constants, so holding onto it would keep a second copy of the model's + weights alive for the whole partitioning pass. """ def __init__( @@ -57,16 +64,23 @@ def __init__( # The builder populates node_info with supported/unsupported status from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder - self._builder = MLXProgramBuilder(edge_program) - self._builder.check_support_only() + builder = MLXProgramBuilder(edge_program) + with mem_phase("partition: check_support"): + builder.check_support_only() + + self._supported: Dict[torch.fx.Node, bool] = {} + self._unsupported_reason: Dict[torch.fx.Node, str] = {} + for node, info in builder.node_info.items(): + self._supported[node] = info.supported + if info.unsupported_reason is not None: + self._unsupported_reason[node] = info.unsupported_reason def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: if node.op != "call_function": return False # Check if builder determined this node is supported - info = self._builder.node_info.get(node) - if info is not None and info.supported: + if self._supported.get(node, False): logger.debug(f"[SUPPORTED] Node {node.target}") return True @@ -86,6 +100,15 @@ def __init__(self, compile_specs: List[CompileSpec] | None = None) -> None: self.compile_specs = compile_specs or [] self.delegation_spec = DelegationSpec(MLXBackend.__name__, self.compile_specs) self.partition_tags: Dict[str, DelegationSpec] = {} + # Last (program, result) pair returned by ops_to_not_decompose(). exir + # asks the same partitioner the same question about the same program + # object twice in a row (_program.py calls it once directly and once + # through _can_skip_using_EDGE_DO_NOT_DECOMP), and answering costs a + # full builder run that repacks every quantized weight. + self._not_decompose_cache: ( + Tuple["weakref.ReferenceType[ExportedProgram]", List[torch._ops.OpOverload]] + | None + ) = None def ops_to_not_decompose( self, ep: ExportedProgram @@ -104,6 +127,12 @@ def ops_to_not_decompose( the shape_env. build() calls _build_mlx_graph() which evaluates SymInts to concrete values when converting tensor shapes, which corrupts the shape_env and causes dynamic shapes to be lost during decomposition. + + Support has to be decided by actually running the handlers: a registered + handler is not proof that a node lowers (aten.layer_norm.default has a + handler that rejects the 6-arg edge form, for instance). Preserving an op + the handler then rejects is worse than not preserving it, because the op + neither decomposes into something delegatable nor lowers itself. """ from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder @@ -116,9 +145,15 @@ def ops_to_not_decompose( ) return ([], None) + cached = self._not_decompose_cache + if cached is not None and cached[0]() is ep: + logger.debug("MLX ops_to_not_decompose: reusing result for same program") + return (cached[1], None) + # Run the builder to determine which nodes are supported builder = MLXProgramBuilder(ep) - builder.check_support_only() + with mem_phase("ops_to_not_decompose: check_support"): + builder.check_support_only() # Collect ops for nodes that are actually supported do_not_decompose: list[torch._ops.OpOverload] = [] @@ -132,6 +167,8 @@ def ops_to_not_decompose( if node.target not in do_not_decompose: do_not_decompose.append(node.target) + self._not_decompose_cache = (weakref.ref(ep), do_not_decompose) + logger.debug( f"MLX ops_to_not_decompose: {[str(op) for op in do_not_decompose]}" ) @@ -152,8 +189,9 @@ def generate_partitions(self, edge_program: ExportedProgram) -> List[Any]: is_supported = self.supported_ops.is_node_supported({}, node) if not is_supported and node.op == "call_function": target_str = str(node.target) - info = self.supported_ops._builder.node_info.get(node) - reason = info.unsupported_reason if info else "No handler registered" + reason = self.supported_ops._unsupported_reason.get( + node, "No handler registered" + ) if target_str in unsupported_by_target: count, _ = unsupported_by_target[target_str] unsupported_by_target[target_str] = (count + 1, reason) diff --git a/backends/mlx/patterns.py b/backends/mlx/patterns.py index 1ceb6064f3b..1760dc63b9a 100644 --- a/backends/mlx/patterns.py +++ b/backends/mlx/patterns.py @@ -22,6 +22,7 @@ emit_quantized_biases, emit_quantized_gather, emit_stop_position, + mlx_qparams_supported, parse_dequant_int4_node, parse_dequant_mx_node, parse_dequant_node, @@ -86,6 +87,44 @@ def _unpack_int4_to_intx_fields( return q, scale_nk, zero_point_nk +def _affine_qparams_lower( + qdata_node: Node, scale_node: Node, group_size: int, bits: int +) -> bool: + """Whether the affine handlers can repack these params, from meta alone. + + ``qdata`` is ``(rows, in_features)`` int8 and ``scale`` is + ``[..., in_features // weight_group_size]``. + """ + qdata = qdata_node.meta.get("val", None) + scale = scale_node.meta.get("val", None) + if qdata is None or scale is None: + return False + if qdata.dim() != 2 or qdata.dtype != torch.int8: + return False + in_features, num_groups = qdata.shape[-1], scale.shape[-1] + if not isinstance(in_features, int) or not isinstance(num_groups, int): + return False + return mlx_qparams_supported(in_features, num_groups, group_size, bits) + + +def _int4_qparams_lower(qdata_node: Node, scale_node: Node, group_size: int) -> bool: + """Same, for the ``Int4Tensor`` layout that _unpack_int4_to_intx_fields reads. + + ``qdata`` is ``(N, K//2)`` nibble-packed single-byte values and ``scale`` is + ``(K // weight_group_size, N)`` -- note the transpose relative to affine. + """ + qdata = qdata_node.meta.get("val", None) + scale = scale_node.meta.get("val", None) + if qdata is None or scale is None: + return False + if qdata.dim() != 2 or qdata.element_size() != 1 or scale.dim() != 2: + return False + packed_cols, num_groups = qdata.shape[-1], scale.shape[0] + if not isinstance(packed_cols, int) or not isinstance(num_groups, int): + return False + return mlx_qparams_supported(2 * packed_cols, num_groups, group_size, 4) + + @REGISTRY.register_pattern(name="INDEX_COPY") class IndexCopyHandler(PatternHandler): """ @@ -892,6 +931,9 @@ def maybe_create( out_dtype=out_dtype, ) + def supported(self, P: MLXProgramBuilder, n: Node) -> bool: + return _affine_qparams_lower(self.qdata, self.scale, self.group_size, self.bits) + def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: assert n == self.head @@ -1022,6 +1064,9 @@ def maybe_create( out_dtype=out_dtype, ) + def supported(self, P: MLXProgramBuilder, n: Node) -> bool: + return _affine_qparams_lower(self.qdata, self.scale, self.group_size, self.bits) + def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: assert n == self.head indices_node = n.args[1] @@ -1396,6 +1441,9 @@ def maybe_create(cls, ep, head): return None return cls(head, [dequant], qdata, scale, zero_point, group_size, out_dtype) + def supported(self, P: MLXProgramBuilder, n: Node) -> bool: + return _int4_qparams_lower(self.qdata, self.scale, self.group_size) + def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: assert n == self.head x_node = n.args[0] @@ -1491,6 +1539,9 @@ def maybe_create(cls, ep, head): qdata, scale, zero_point, group_size, out_dtype = parsed return cls(head, [dequant], qdata, scale, zero_point, group_size, out_dtype) + def supported(self, P: MLXProgramBuilder, n: Node) -> bool: + return _int4_qparams_lower(self.qdata, self.scale, self.group_size) + def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: assert n == self.head indices_node = n.args[1] diff --git a/backends/mlx/preprocess.py b/backends/mlx/preprocess.py index 315835f1689..b2e26a4d2b2 100644 --- a/backends/mlx/preprocess.py +++ b/backends/mlx/preprocess.py @@ -22,6 +22,7 @@ from typing import ClassVar, final, List from executorch.backends.mlx._logging import logger +from executorch.backends.mlx._memprofile import log_footprint, mem_phase from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder from executorch.backends.mlx.serialization.mlx_graph_serialize import ( HEADER_LENGTH, @@ -74,6 +75,7 @@ def preprocess( """ logger.debug("MLXBackend.preprocess() called") logger.debug(f"Edge program:\n{edge_program}") + log_footprint("preprocess entry") # Build MLXGraph from ExportedProgram # Use a deterministic 4-hex prefix derived from the edge program to @@ -82,22 +84,32 @@ def preprocess( # with the same auto-generated name. prefix = hashlib.sha256(str(edge_program).encode()).hexdigest()[:4] builder = MLXProgramBuilder(edge_program, named_data_key_prefix=prefix) - mlx_graph = builder.build() - # Get constant data as NamedDataStore (ET will own this data) - named_data_store = builder.get_named_data_store() + with mem_phase(f"preprocess[{prefix}]"): + with mem_phase("build"): + mlx_graph = builder.build() - logger.debug(f" named_data_store entries: {len(named_data_store.pte_data)}") - _log_mlx_graph(mlx_graph) + # Get constant data as NamedDataStore (ET will own this data) + with mem_phase("get_named_data_store"): + named_data_store = builder.get_named_data_store() - # Serialize to bytes (no constant data embedded) - serialized = serialize_mlx_graph(mlx_graph) + logger.debug( + f" named_data_store entries: {len(named_data_store.pte_data)}" + ) + _log_mlx_graph(mlx_graph) + + # Serialize to bytes (no constant data embedded) + with mem_phase("serialize"): + serialized = serialize_mlx_graph(mlx_graph) + + with mem_phase("data_store_output"): + data_store_output = named_data_store.get_named_data_store_output() logger.debug(f"MLXBackend.preprocess() complete: {len(serialized)} bytes") return PreprocessResult( processed_bytes=serialized, - data_store_output=named_data_store.get_named_data_store_output(), + data_store_output=data_store_output, ) diff --git a/examples/models/muse-glimmer/CMakePresets.json b/examples/models/muse-glimmer/CMakePresets.json index ab186e6001d..4cc407d4081 100644 --- a/examples/models/muse-glimmer/CMakePresets.json +++ b/examples/models/muse-glimmer/CMakePresets.json @@ -28,7 +28,9 @@ "name": "muse-glimmer-mlx", "displayName": "Muse Glimmer runner (MLX)", "inherits": ["muse-glimmer-base"], - "cacheVariables": {}, + "cacheVariables": { + "EXECUTORCH_BUILD_MLX": "ON" + }, "condition": { "type": "equals", "lhs": "${hostSystemName}", diff --git a/examples/models/muse-glimmer/README.md b/examples/models/muse-glimmer/README.md index 98e91eed37c..657cfd716b1 100644 --- a/examples/models/muse-glimmer/README.md +++ b/examples/models/muse-glimmer/README.md @@ -126,7 +126,19 @@ Vision adds `vision_encoder` to each method set. ## Build the runners -Use the model-specific CMake workflow preset for the selected backend: +From the repo root, the make targets build ExecuTorch with the backend and then +the runners: + +```bash +# CUDA +make muse-glimmer-cuda + +# MLX +make muse-glimmer-mlx +``` + +When ExecuTorch is already built for that backend, the model's CMake workflow +preset rebuilds just the runners: ```bash # CUDA