Update

[ghstack-poisoned]

Author: JCNTH

backends/arm/tosa/specification.py

@@ -13,10 +13,6 @@
 import re
 from typing import Dict, Generic, List, Set, TypeVar
 
-# Import torch before any torch submodule so torch.SymInt is defined; importing
-# torch.fx first triggers a torch.distributed circular import during collection.
-import torch  # noqa: F401
-
 from packaging.version import Version
 from torch.fx.experimental.symbolic_shapes import ShapeEnv
 

backends/mlx/CMakeLists.txt

@@ -162,78 +162,83 @@ if(NOT _mlx_deployment_target_ok)
   )
 endif()
 
-# Build MLX in its own isolated CMake scope via ExternalProject, then consume it
-# as a prebuilt static lib through an imported target (see below). Building MLX
-# with add_subdirectory would drop its whole project into ExecuTorch's
-# target/option namespace, which collides with shared deps MLX fetches (e.g.
-# nlohmann_json) and leaks MLX's MLX_BUILD_* options into our cache. The
-# isolated scope runs MLX's FetchContent in its own namespace, so no collision
-# and no submodule patching are needed.
-include(ExternalProject)
-
-set(_mlx_binary_dir ${CMAKE_CURRENT_BINARY_DIR}/mlx)
-set(_mlx_static_lib ${_mlx_binary_dir}/libmlx.a)
-# With MLX_METAL_JIT=ON, MLX does not install the metallib; it is produced in
-# the build tree at this path.
-set(_mlx_metallib ${_mlx_binary_dir}/mlx/backend/metal/kernels/mlx.metallib)
-
-message(
-  STATUS "Building MLX from submodule (ExternalProject): ${MLX_SOURCE_DIR}"
+# MLX build options - we only need the C++ library with Metal
+set(MLX_BUILD_PYTHON_BINDINGS
+    OFF
+    CACHE BOOL "" FORCE
 )
-ExternalProject_Add(
-  mlx_external
-  SOURCE_DIR ${MLX_SOURCE_DIR}
-  BINARY_DIR ${_mlx_binary_dir}
-  CMAKE_ARGS -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
-             -DCMAKE_CXX_STANDARD=${CMAKE_CXX_STANDARD}
-             -DCMAKE_OSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}
-             -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TOOLCHAIN_FILE}
-             -DPLATFORM=${PLATFORM}
-             -DDEPLOYMENT_TARGET=${DEPLOYMENT_TARGET}
-             # We only need the static Metal C++ library.
-             -DMLX_BUILD_METAL=ON
-             -DMLX_BUILD_CPU=OFF
-             -DMLX_BUILD_CUDA=OFF
-             -DMLX_BUILD_SHARED_LIBS=OFF
-             -DMLX_BUILD_PYTHON_BINDINGS=OFF
-             -DMLX_BUILD_PYTHON_STUBS=OFF
-             -DMLX_BUILD_TESTS=OFF
-             -DMLX_BUILD_EXAMPLES=OFF
-             -DMLX_BUILD_BENCHMARKS=OFF
-             -DMLX_BUILD_GGUF=OFF
-             -DMLX_BUILD_SAFETENSORS=OFF
-             -DMLX_METAL_JIT=ON
-  # MLX's own install() does not emit libmlx.a where we consume it or the
-  # metallib at all, so skip the install step and read both from the build tree.
-  INSTALL_COMMAND ""
-  # ExternalProject stamps its build, so a bare MLX submodule bump (git
-  # submodule update) would not invalidate the stamp and we'd link a stale
-  # libmlx.a with no signal. BUILD_ALWAYS reruns the build step every configure;
-  # it is a fast no-op under Ninja when nothing changed.
-  BUILD_ALWAYS ON
-  # Required so Ninja knows these are produced by the external build.
-  BUILD_BYPRODUCTS ${_mlx_static_lib} ${_mlx_metallib}
+set(MLX_BUILD_TESTS
+    OFF
+    CACHE BOOL "" FORCE
 )
-
-# Imported target for the MLX static library produced by mlx_external. A static
-# libmlx.a carries no transitive link deps, so re-add the frameworks MLX itself
-# links (mirrors third-party/mlx/CMakeLists.txt:209). CPU is OFF, so Accelerate
-# is not needed.
-add_library(mlx STATIC IMPORTED GLOBAL)
-set_target_properties(mlx PROPERTIES IMPORTED_LOCATION ${_mlx_static_lib})
-# Headers come from the submodule source tree (always present), not the build
-# dir, so mlxdelegate compilation cannot race the external build. Verified: MLX
-# emits no public headers into the build tree (version.h is checked in; the
-# version string is a compile-def on version.cpp), so the source dir suffices.
-target_include_directories(mlx INTERFACE ${MLX_SOURCE_DIR})
-find_library(METAL_FRAMEWORK Metal)
-find_library(FOUNDATION_FRAMEWORK Foundation)
-find_library(QUARTZ_FRAMEWORK QuartzCore)
-set_target_properties(
-  mlx
-  PROPERTIES INTERFACE_LINK_LIBRARIES
-             "${METAL_FRAMEWORK};${FOUNDATION_FRAMEWORK};${QUARTZ_FRAMEWORK}"
+set(MLX_BUILD_EXAMPLES
+    OFF
+    CACHE BOOL "" FORCE
+)
+set(MLX_BUILD_BENCHMARKS
+    OFF
+    CACHE BOOL "" FORCE
+)
+set(MLX_BUILD_PYTHON_STUBS
+    OFF
+    CACHE BOOL "" FORCE
+)
+set(MLX_BUILD_CUDA
+    OFF
+    CACHE BOOL "" FORCE
+)
+set(MLX_BUILD_CPU
+    OFF
+    CACHE BOOL "" FORCE
+)
+set(MLX_BUILD_METAL
+    ON
+    CACHE BOOL "" FORCE
+)
+set(MLX_BUILD_SHARED_LIBS
+    OFF
+    CACHE BOOL "" FORCE
+)
+set(MLX_BUILD_GGUF
+    OFF
+    CACHE BOOL "" FORCE
+)
+set(MLX_BUILD_SAFETENSORS
+    OFF
+    CACHE BOOL "" FORCE
 )
+set(MLX_METAL_JIT
+    ON
+    CACHE BOOL "Use JIT compiled Metal kernels"
+)
+
+# Auto-apply patches to MLX submodule. Each patch is applied idempotently: `git
+# apply --check` tests whether the patch is still applicable (i.e. not yet
+# applied), and only then applies it.
+set(_mlx_patches "${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_json.patch")
+foreach(_patch IN LISTS _mlx_patches)
+  if(EXISTS "${_patch}" AND EXISTS "${MLX_SOURCE_DIR}")
+    get_filename_component(_patch_name "${_patch}" NAME)
+    execute_process(
+      COMMAND git apply --check "${_patch}"
+      WORKING_DIRECTORY ${MLX_SOURCE_DIR}
+      RESULT_VARIABLE _patch_check
+      OUTPUT_QUIET ERROR_QUIET
+    )
+    if(_patch_check EQUAL 0)
+      execute_process(
+        COMMAND git apply "${_patch}" WORKING_DIRECTORY ${MLX_SOURCE_DIR}
+      )
+      message(STATUS "Applied ${_patch_name} to MLX submodule")
+    else()
+      message(STATUS "${_patch_name} already applied or not applicable")
+    endif()
+  endif()
+endforeach()
+
+# Add MLX subdirectory
+message(STATUS "Adding MLX from submodule: ${MLX_SOURCE_DIR}")
+add_subdirectory(${MLX_SOURCE_DIR} ${CMAKE_CURRENT_BINARY_DIR}/mlx)
 
 # -----------------------------------------------------------------------------
 # MLX Backend library
@@ -261,10 +266,6 @@ add_library(mlxdelegate ${_mlx_backend__srcs})
 # Ensure schema is generated before compiling
 add_dependencies(mlxdelegate mlx_schema)
 
-# mlx is an imported target, so a dependency on it would not order the build.
-# Depend on mlx_external directly so libmlx.a exists before mlxdelegate links.
-add_dependencies(mlxdelegate mlx_external)
-
 # Add logging flag if enabled
 if(ET_MLX_ENABLE_OP_LOGGING)
   target_compile_definitions(mlxdelegate PRIVATE ET_MLX_ENABLE_OP_LOGGING=1)
@@ -282,9 +283,8 @@ target_include_directories(
   mlxdelegate PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/runtime
 )
 
-# Link against MLX and executorch. mlx is an imported static lib (built by
-# mlx_external); it is only consumed at build time here, not re-exported, so
-# downstream consumers must link to mlx separately via the package config.
+# Link against MLX and executorch mlx is only available at BUILD_INTERFACE -
+# consumers must link to mlx separately
 target_link_libraries(
   mlxdelegate PRIVATE mlx_schema executorch_core $<BUILD_INTERFACE:mlx>
 )
@@ -301,10 +301,8 @@ install(
   DESTINATION ${CMAKE_INSTALL_LIBDIR}
 )
 
-# Install mlx library for downstream consumers. mlx is an imported target (built
-# by mlx_external), so install the static lib file directly rather than via
-# install(TARGETS ...), which only accepts built targets.
-install(FILES ${_mlx_static_lib} DESTINATION ${CMAKE_INSTALL_LIBDIR})
+# Install mlx library for downstream consumers
+install(TARGETS mlx DESTINATION ${CMAKE_INSTALL_LIBDIR})
 
 # Install mlx headers for downstream consumers that may need mlx types
 install(
@@ -327,15 +325,16 @@ install(
 # containing MLX code. When MLX is statically linked into _portable_lib.so, this
 # is the directory containing _portable_lib.so.
 #
-# For the installed library, we put metallib in lib/ alongside libmlx.a. The
-# metallib is produced in the mlx_external build tree (MLX_METAL_JIT=ON does not
-# install it); _mlx_metallib points there.
-install(FILES ${_mlx_metallib} DESTINATION ${CMAKE_INSTALL_LIBDIR})
+# For the installed library, we put metallib in lib/ alongside libmlx.a
+install(
+  FILES ${CMAKE_CURRENT_BINARY_DIR}/mlx/mlx/backend/metal/kernels/mlx.metallib
+  DESTINATION ${CMAKE_INSTALL_LIBDIR}
+)
 
 # Cache the metallib path for pybindings to copy it next to _portable_lib.so
 # This enables editable installs to work correctly
 set(MLX_METALLIB_PATH
-    "${_mlx_metallib}"
+    "${CMAKE_CURRENT_BINARY_DIR}/mlx/mlx/backend/metal/kernels/mlx.metallib"
     CACHE INTERNAL "Path to mlx.metallib for pybindings"
 )
 

backends/mlx/llm/sampling.py

@@ -12,8 +12,7 @@
 
 class SamplingHead(nn.Module):
     """
-    Wraps a model that returns last-token logits ``(B, vocab)`` and samples a
-    token id ``(B)`` on-device.
+    Wraps a model that returns logits and samples a token id on-device.
 
         forward(*model_args, temperature, top_k=None, top_p=1.0, seed=None,
                 **model_kwargs) -> token_id
@@ -34,11 +33,12 @@ def __init__(self, model: nn.Module):
         self.model = model
 
     def forward(self, *args, temperature, top_k=None, top_p=1.0, seed=None, **kwargs):
-        logits = self.model(*args, **kwargs)  # [B, vocab]
+        logits = self.model(*args, **kwargs)  # [B, S, vocab]
+        last = logits[:, -1, :]  # [B, vocab]
         if not isinstance(top_p, torch.Tensor):
             top_p = torch.tensor(float(top_p))
         if top_k is None:
             top_k = torch.tensor(torch.iinfo(torch.int64).max, dtype=torch.int64)
         elif not isinstance(top_k, torch.Tensor):
             top_k = torch.tensor(int(top_k), dtype=torch.int64)
-        return torch.ops.mlx.sample(logits, temperature, top_k, top_p, seed)
+        return torch.ops.mlx.sample(last, temperature, top_k, top_p, seed)

backends/mlx/patches/mlx_json.patch

@@ -0,0 +1,29 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -304,12 +304,18 @@ else()
+   set(MLX_BUILD_ACCELERATE OFF)
+ endif()
+
+-message(STATUS "Downloading json")
+-FetchContent_Declare(
+-  json
+-  URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz)
+-FetchContent_MakeAvailable(json)
+-target_include_directories(
+-  mlx PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
++# Only fetch json if nlohmann_json target doesn't already exist
++# (ExecuTorch provides its own copy)
++if(NOT TARGET nlohmann_json)
++  message(STATUS "Downloading json")
++  FetchContent_Declare(
++    json
++    URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz)
++  FetchContent_MakeAvailable(json)
++  target_include_directories(
++    mlx PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
++else()
++  message(STATUS "Using existing nlohmann_json target")
++endif()
+
+ add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)

backends/mlx/runtime/MLXBackend.cpp

@@ -19,8 +19,6 @@
 
 #include <mlx/mlx.h>
 
-#include <executorch/backends/mlx/runtime/backend_options.h>
-
 #include <cstring>
 #include <limits>
 #include <memory>
@@ -171,12 +169,6 @@ struct MLXHandle {
   Interpreter interpreter;
   ::mlx::core::Stream stream; // Dedicated GPU stream for this handle
 
-  // Periodically call mlx::core::clear_cache() every N execute() calls to
-  // release MLX's cached buffer pool. 0 disables. Counter is non-atomic: it is
-  // only touched inside mlx_global_mutex() (see execute()).
-  int clear_cache_interval_{0};
-  uint64_t execute_count_{0};
-
   // Keep the constant buffers alive for zero-copy constants
   // Each FreeableBuffer must outlive the MLX arrays that reference it
   std::vector<FreeableBuffer> constant_buffers;
@@ -219,13 +211,6 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
     try {
       new (handle) MLXHandle();
 
-      // Per-model clear-cache cadence (optional runtime spec, keyed per
-      // delegate)
-      if (auto spec = context.get_runtime_spec<int>(kClearCacheIntervalKey);
-          spec.ok() && spec.get() > 0) {
-        handle->clear_cache_interval_ = spec.get();
-      }
-
       if (!processed || !processed->data() || processed->size() == 0) {
         throw std::runtime_error("init: null or empty delegate payload");
       }
@@ -269,42 +254,8 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
       processed->Free();
       processed = nullptr;
 
-      // Load mutable buffers (e.g., KV cache). When skip_mutable_buffer_init is
-      // set (multi-session, where mlx_mutable_state.h allocates per-session
-      // buffers), avoid keeping a dead default copy on the handle. This is only
-      // safe if the init chain does not reference mutable buffers — otherwise
-      // it would run against an empty store, so we error out clearly.
-      bool skip_mutable_buffer_init = false;
-      if (auto spec = context.get_runtime_spec<bool>(kSkipMutableBufferInitKey);
-          spec.ok()) {
-        skip_mutable_buffer_init = spec.get();
-      }
-      // skip_mutable_buffer_init is only safe under a multi-session owner: the
-      // per-session manager allocates the buffers and execute() rebinds to
-      // them. Without an active load scope no handle is registered, no
-      // per-session buffers are ever created, and execute() would run against
-      // empty buffers. Reject the misuse loudly here instead of failing later.
-      if (skip_mutable_buffer_init && !mutable_state_load_scope_active()) {
-        ET_LOG(
-            Error,
-            "skip_mutable_buffer_init set without an active multi-session load "
-            "scope; mutable buffers would never be allocated");
-        throw std::runtime_error(
-            "skip_mutable_buffer_init requires an active multi-session owner");
-      }
-      if (skip_mutable_buffer_init &&
-          init_chain_references_mutable_buffer(handle->program)) {
-        ET_LOG(
-            Error,
-            "skip_mutable_buffer_init set but the init chain references "
-            "mutable buffers; cannot skip default mutable-buffer init");
-        throw std::runtime_error(
-            "skip_mutable_buffer_init incompatible with init chain that "
-            "references mutable buffers");
-      }
-      if (!skip_mutable_buffer_init) {
-        load_mutable_buffers(handle->program, handle->mutable_buffers);
-      }
+      // Load mutable buffers (e.g., KV cache)
+      load_mutable_buffers(handle->program, handle->mutable_buffers);
 
       // Bind execution state (reused across execute() calls)
       handle->state.bind(
@@ -482,16 +433,6 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
 
       h->state.reset(); // Release temp GPU buffers back to MLX cache
 
-      // Periodically release MLX's cached buffer pool. The counter increment
-      // and clear_cache() run under the global mutex so they can't race another
-      // handle's in-flight submission, and the counter needs no atomic.
-      if (h->clear_cache_interval_ > 0) {
-        std::lock_guard<std::mutex> lock(mlx_global_mutex());
-        if ((++h->execute_count_ % h->clear_cache_interval_) == 0) {
-          ::mlx::core::clear_cache();
-        }
-      }
-
       return Error::Ok;
     } catch (const std::exception& e) {
       ET_LOG(Error, "MLX execute failed: %s", e.what());
@@ -514,7 +455,7 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
 
 namespace {
 auto cls = MLXBackend();
-Backend backend{kMLXBackendId, &cls};
+Backend backend{"MLXBackend", &cls};
 static auto success_with_compiler = register_backend(backend);
 } // namespace