up

Author: metascroy

.github/workflows/mlx.yml

@@ -95,7 +95,7 @@ jobs:
         fi
 
         if [ "${{ matrix.suite }}" = "operators" ]; then
-          MAX_FAILURES=1
+          MAX_FAILURES=0
         else
           MAX_FAILURES=3
         fi

backends/mlx/runtime/MLXBackend.cpp

@@ -390,6 +390,8 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
         write_output(prepared_outputs[info.prepared_idx], out_tensor);
       }
 
+      h->state.reset(); // Release temp GPU buffers back to MLX cache
+
       return Error::Ok;
     } catch (const std::exception& e) {
       ET_LOG(Error, "MLX execute failed: %s", e.what());

backends/mlx/runtime/MLXInterpreter.h

@@ -924,6 +924,9 @@ inline void exec_slice_update(
     ExecutionState& st,
     StreamOrDevice s) {
   // When out == dst, use direct assignment to preserve MLX buffer donation.
+  // TODO: I'm not sure if this is needed as a special case since the standard
+  // st.set_tensor does a std::move. Keeping for now, but should investigate and
+  // possibly remove in future.
   const bool in_place = (n.out.idx == n.dst.idx);
   array& dst = st.tensor_ref(n.dst);
   const array& upd = st.const_tensor_ref(n.update);
@@ -1060,6 +1063,10 @@ exec_index_copy(const IndexCopyNode& n, ExecutionState& st, StreamOrDevice s) {
     idx_vec[i] = static_cast<int>(idx);
   }
 
+  // When out == dst, use direct assignment to preserve MLX buffer donation.
+  // TODO: I'm not sure if this is needed as a special case since the standard
+  // st.set_tensor does a std::move. Keeping for now, but should investigate and
+  // possibly remove in future.
   const bool in_place = (n.out.idx == n.dst.idx);
 
   if (idx_vec.empty()) {
@@ -1995,8 +2002,6 @@ class Interpreter {
       case OpCode::ARG_PARTITION:
         ops::exec_argpartition(std::get<ArgPartitionNode>(instr.node), st, s);
         break;
-      case OpCode::SENTINEL:
-        break;
       default:
         throw std::runtime_error(
             "Unknown opcode: " + std::to_string(static_cast<int>(instr.op)));

backends/mlx/serialization/MLXLoader.h.tmpl

@@ -79,15 +79,12 @@ struct IntOrVidOrTid {
 
 enum class OpCode : uint8_t {
 {{OPCODE_ENUM_VALUES}}
-  SENTINEL
 };
 
 // OpCode to string conversion (for logging)
 inline const char* op_name(OpCode op) {
   switch (op) {
 {{OP_NAME_CASES}}
-    case OpCode::SENTINEL:
-      return "SENTINEL";
   }
   return "UNKNOWN";
 }

extension/llm/export/config/llm_config.py

@@ -558,6 +558,7 @@ class EthosUConfig:
     system_config: str = "default"
 
 
+@dataclass
 class MLXConfig:
     """
     Configures the MLX backend for Apple Silicon.