Update

[ghstack-poisoned]

Author: JCNTH

backends/arm/tosa/specification.py

@@ -13,6 +13,10 @@
 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/llm/sampling.py

@@ -12,7 +12,8 @@
 
 class SamplingHead(nn.Module):
     """
-    Wraps a model that returns logits and samples a token id on-device.
+    Wraps a model that returns last-token logits ``(B, vocab)`` and samples a
+    token id ``(B)`` on-device.
 
         forward(*model_args, temperature, top_k=None, top_p=1.0, seed=None,
                 **model_kwargs) -> token_id
@@ -33,12 +34,11 @@ 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, S, vocab]
-        last = logits[:, -1, :]  # [B, vocab]
+        logits = self.model(*args, **kwargs)  # [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(last, temperature, top_k, top_p, seed)
+        return torch.ops.mlx.sample(logits, temperature, top_k, top_p, seed)

backends/mlx/runtime/MLXBackend.cpp

@@ -19,6 +19,8 @@
 
 #include <mlx/mlx.h>
 
+#include <executorch/backends/mlx/runtime/backend_options.h>
+
 #include <cstring>
 #include <limits>
 #include <memory>
@@ -169,6 +171,12 @@ 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;
@@ -211,6 +219,13 @@ 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");
       }
@@ -254,8 +269,42 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
       processed->Free();
       processed = nullptr;
 
-      // Load mutable buffers (e.g., KV cache)
-      load_mutable_buffers(handle->program, handle->mutable_buffers);
+      // 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);
+      }
 
       // Bind execution state (reused across execute() calls)
       handle->state.bind(
@@ -433,6 +482,16 @@ 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());
@@ -455,7 +514,7 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
 
 namespace {
 auto cls = MLXBackend();
-Backend backend{"MLXBackend", &cls};
+Backend backend{kMLXBackendId, &cls};
 static auto success_with_compiler = register_backend(backend);
 } // namespace
 

backends/mlx/runtime/backend_options.h

@@ -0,0 +1,47 @@
+/*
+ * 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 option keys for the MLX backend. Included by the backend itself
+// (to read the per-model runtime spec) and by callers/runners (to set it via
+// a LoadBackendOptionsMap), keeping the string literals in one place.
+
+#pragma once
+
+namespace executorch {
+namespace backends {
+namespace mlx {
+
+// Backend id under which the MLX backend registers (see MLXBackend.cpp).
+inline constexpr char kMLXBackendId[] = "MLXBackend";
+
+// Per-model runtime-spec key. Value N means: call mlx::core::clear_cache()
+// every N execute() calls to release MLX's cached buffer pool; 0/unset
+// disables.
+//
+// NOTE on granularity: MLX's buffer cache is a process-global singleton, so the
+// flush is global even though this key is read per delegate handle. The counter
+// is per-handle, so with a single MLX handle (the common case — gemma's
+// prefill/decode share one "forward" method) the cadence is exactly "every N
+// forwards"; if a process loads multiple MLX handles, the effective cadence is
+// the aggregate of their executes and any handle's flush frees the shared pool
+// for all. This bounds resident-*average*: between flushes the cache can still
+// grow to MLX's default ceiling, and each flush is followed by a cold-cache
+// realloc. A future set_cache_limit-style key could complement this by bounding
+// peak footprint continuously.
+inline constexpr char kClearCacheIntervalKey[] = "clear_cache_interval";
+
+// Per-model runtime-spec key (bool). When true, the handle does not allocate
+// its own default mutable buffers at init() — per-session buffers are managed
+// by mlx_mutable_state.h instead. Only valid for multi-session loads, and only
+// when the program's init chain does not reference mutable buffers (init()
+// errors otherwise). Saves one full mutable-buffer (KV-cache) copy per handle.
+inline constexpr char kSkipMutableBufferInitKey[] = "skip_mutable_buffer_init";
+
+} // namespace mlx
+} // namespace backends
+} // namespace executorch

backends/mlx/runtime/mlx_mutable_state.cpp

@@ -136,15 +136,52 @@ int64_t mutable_state_bytes_per_session(MutableStateContext ctx) {
     return 0;
   }
   int64_t total = 0;
+  // Per-session mutable buffers are allocated from program metadata
+  // (mutable_buffer_map + tensor_meta), independent of whether the handle kept
+  // a default copy. Compute the estimate from metadata so it stays correct even
+  // when default mutable-buffer init was skipped (skip_mutable_buffer_init).
   for (const auto& kv : it->second.handles) {
-    const MutableBufferData* bufs = kv.second.default_buffers;
-    if (bufs == nullptr) {
+    const MLXProgram* program = kv.second.program;
+    if (program == nullptr) {
       continue;
     }
-    for (const auto& t : bufs->tensors) {
-      if (t.has_value()) {
-        total += static_cast<int64_t>(t->nbytes());
+    for (const auto& slot : program->mutable_buffer_map) {
+      if (slot.slot_type != SlotType::TensorSlot ||
+          slot.idx >= program->tensor_meta.size() ||
+          !program->tensor_meta[slot.idx].has_value()) {
+        continue;
       }
+      const auto& meta = *program->tensor_meta[slot.idx];
+      // Sum sizes from metadata, clamping each tensor to kMaxAllocationBytes so
+      // malformed (oversized) shapes in an untrusted program can't overflow the
+      // accumulator. Real allocations are independently bounded by
+      // check_allocation_bounded at load_mutable_buffers.
+      uint64_t bytes = static_cast<uint64_t>(
+          ::mlx::core::size_of(resolve_dtype(meta.scalar_type)));
+      bool dynamic = false;
+      for (const auto& dim : meta.shape) {
+        if (dim.value < 0) {
+          dynamic = true;
+          break;
+        }
+        const uint64_t d = static_cast<uint64_t>(dim.value);
+        if (d == 0) {
+          bytes = 0;
+          break;
+        }
+        if (bytes > kMaxAllocationBytes / d) {
+          bytes = kMaxAllocationBytes; // clamp; avoids overflow
+        } else {
+          bytes *= d;
+        }
+      }
+      if (dynamic) {
+        continue;
+      }
+      if (bytes > kMaxAllocationBytes) {
+        bytes = kMaxAllocationBytes;
+      }
+      total += static_cast<int64_t>(bytes);
     }
   }
   return total;
@@ -217,6 +254,10 @@ void mutable_state_note_handle(
   handle_ctx()[handle] = tl_loading_ctx;
 }
 
+bool mutable_state_load_scope_active() {
+  return tl_loading_ctx != kInvalidMutableContext;
+}
+
 void mutable_state_forget_handle(const void* handle) {
   std::lock_guard<std::mutex> g(registry_mutex());
   auto hit = handle_ctx().find(handle);

backends/mlx/runtime/mlx_mutable_state.h

@@ -195,6 +195,12 @@ void mutable_state_note_handle(
 
 void mutable_state_forget_handle(const void* handle);
 
+// True if a multi-session load scope (with_load_scope) is active on the current
+// thread. The MLXBackend uses this to reject skip_mutable_buffer_init when no
+// owner is active for the load — without an owner, no per-session buffers are
+// ever allocated and execute() would run against empty mutable buffers.
+bool mutable_state_load_scope_active();
+
 ::executorch::runtime::Error mutable_state_rebind_for_execute(
     const void* handle,
     ExecutionState& state);

backends/mlx/serialization/MLXLoader.h.tmpl

@@ -118,6 +118,25 @@ struct Instruction {
   }
 };
 
+// =============================================================================
+// for_each_tid - invokes cb(Tid) for every tensor id referenced by an
+// instruction's node (AUTO-GENERATED from schema.fbs).
+//
+// NOTE: nested chain-index fields (ScanNode::body_chain_idx,
+// IfNode::then_chain_idx/else_chain_idx) are NOT followed here; callers that
+// need transitive coverage must recurse into those chains themselves.
+// =============================================================================
+
+template <typename F>
+inline void for_each_tid(const Instruction& instr, F&& cb) {
+  switch (instr.op) {
+{{FOR_EACH_TID_CASES}}
+    default:
+      throw std::runtime_error(
+          std::string("for_each_tid: unhandled op ") + op_name(instr.op));
+  }
+}
+
 // =============================================================================
 // SlotVariant for I/O mapping
 // =============================================================================
@@ -194,6 +213,65 @@ struct MLXProgram {
   }
 };
 
+// =============================================================================
+// init_chain_references_mutable_buffer - returns true if the program's init
+// chain references any mutable-buffer Tid, following nested SCAN/IF chains.
+//
+// Used to decide whether the handle's default mutable-buffer initialization can
+// be safely skipped (e.g. when buffers are managed per-session). If the init
+// chain reads/writes mutable buffers, skipping would be unsafe.
+// =============================================================================
+inline bool init_chain_references_mutable_buffer(const MLXProgram& program) {
+  if (program.init_chain_idx < 0 || program.num_mutable_buffer_tensors == 0) {
+    return false;
+  }
+  // Tid layout: Constant -> Input -> Output -> MutableBuffer -> Temp.
+  const uint32_t output_end = program.num_constant_tensors +
+      program.num_input_tensors + program.num_output_tensors;
+  const uint32_t mutable_buffer_end =
+      output_end + program.num_mutable_buffer_tensors;
+  auto is_mutable_buffer = [&](Tid t) {
+    return t.idx >= output_end && t.idx < mutable_buffer_end;
+  };
+
+  const size_t num_chains = program.instruction_chains.size();
+  std::vector<bool> visited(num_chains, false);
+  std::vector<uint32_t> worklist;
+  worklist.push_back(static_cast<uint32_t>(program.init_chain_idx));
+
+  while (!worklist.empty()) {
+    const uint32_t chain_idx = worklist.back();
+    worklist.pop_back();
+    if (chain_idx >= num_chains || visited[chain_idx]) {
+      continue;
+    }
+    visited[chain_idx] = true;
+    for (const auto& instr : program.instruction_chains[chain_idx]) {
+      bool hit = false;
+      for_each_tid(instr, [&](Tid t) {
+        if (is_mutable_buffer(t)) {
+          hit = true;
+        }
+      });
+      if (hit) {
+        return true;
+      }
+      // Follow nested chains (not covered by for_each_tid).
+      if (instr.op == OpCode::SCAN) {
+        const auto& n = std::get<ScanNode>(instr.node);
+        if (n.body_chain_idx >= 0) {
+          worklist.push_back(static_cast<uint32_t>(n.body_chain_idx));
+        }
+      } else if (instr.op == OpCode::IF) {
+        const auto& n = std::get<IfNode>(instr.node);
+        worklist.push_back(n.then_chain_idx);
+        worklist.push_back(n.else_chain_idx);
+      }
+    }
+  }
+  return false;
+}
+
 // =============================================================================
 // FlatBuffer loading functions
 // =============================================================================