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/repack_mlx.py b/backends/mlx/custom_kernel_ops/gguf/q4k/repack_mlx.py index 16c8543d098..c7114f0756f 100644 --- a/backends/mlx/custom_kernel_ops/gguf/q4k/repack_mlx.py +++ b/backends/mlx/custom_kernel_ops/gguf/q4k/repack_mlx.py @@ -41,13 +41,25 @@ def repack_mlx( ``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). + + The raw blob is released as soon as it has been unpacked, so it does not sit + alongside the packed constants for the rest of the build. """ from executorch.extension.llm.export.gguf import ExportableGGUFTensor - weight_target, raw = P.get_placeholder_target_and_tensor(weight_node) + 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, "q4_k").to_intx_unpacked_to_int8_tensor( max_group_size=128, scale_dtype=scale_dtype ) + del raw + P.release_placeholder_tensor(weight_node) 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 @@ -57,4 +69,6 @@ def repack_mlx( 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 + result = (packed_slot, scales_slot, biases_slot, group_size) + P.repack_cache[weight_target] = result + return result diff --git a/backends/mlx/custom_kernel_ops/gguf/q5k/repack_mlx.py b/backends/mlx/custom_kernel_ops/gguf/q5k/repack_mlx.py index 7477fc68da0..e1fee8e6704 100644 --- a/backends/mlx/custom_kernel_ops/gguf/q5k/repack_mlx.py +++ b/backends/mlx/custom_kernel_ops/gguf/q5k/repack_mlx.py @@ -41,13 +41,25 @@ def repack_mlx( ``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). + + The raw blob is released as soon as it has been unpacked, so it does not sit + alongside the packed constants for the rest of the build. """ from executorch.extension.llm.export.gguf import ExportableGGUFTensor - weight_target, raw = P.get_placeholder_target_and_tensor(weight_node) + 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, "q5_k").to_intx_unpacked_to_int8_tensor( max_group_size=128, scale_dtype=scale_dtype ) + del raw + P.release_placeholder_tensor(weight_node) 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 @@ -57,4 +69,6 @@ def repack_mlx( 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 + result = (packed_slot, scales_slot, biases_slot, group_size) + P.repack_cache[weight_target] = result + return result diff --git a/backends/mlx/custom_kernel_ops/gguf/q6k/repack_mlx.py b/backends/mlx/custom_kernel_ops/gguf/q6k/repack_mlx.py index c2372e338a2..0ac191e4d32 100644 --- a/backends/mlx/custom_kernel_ops/gguf/q6k/repack_mlx.py +++ b/backends/mlx/custom_kernel_ops/gguf/q6k/repack_mlx.py @@ -55,17 +55,31 @@ def repack_mlx( ``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, raw = P.get_placeholder_target_and_tensor(weight_node) + 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, "q6_k").to_intx_unpacked_to_int8_tensor( max_group_size=128, scale_dtype=scale_dtype ) + del raw group_size = int(intx.block_size[-1]) if group_size < _MIN_MLX_GROUP_SIZE: + # Falling back to the fused kernels, which consume the raw blob. return None + 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, _BITS) @@ -84,4 +98,6 @@ def repack_mlx( biases, scales_slot, ) - return packed_slot, scales_slot, biases_slot, group_size + 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, )