mlx-bun 0.0.12 → 0.2.0

package/src/web/app.js

@@ -3496,15 +3496,21 @@ function createChatController() {
       t0: performance.now(),
       tFirst: 0,
       tokens: 0,
+      settled: false,
       blockState: { blocks: [] },
       scheduleText: () => {},
       scheduleThinking: () => {},
       citations: pendingCitations
     };
     pendingCitations = [];
-    curAssistant.scheduleText = makeFrameScheduler(() => renderBlocksIncremental(curAssistant.textNode, curAssistant.text, curAssistant.blockState), atBottom, stick);
-    curAssistant.scheduleThinking = makeFrameScheduler(() => {
-      curAssistant.thinkBody.textContent = curAssistant.thinking;
+    const assistant = curAssistant;
+    assistant.scheduleText = makeFrameScheduler(() => {
+      if (!assistant.settled)
+        renderBlocksIncremental(assistant.textNode, assistant.text, assistant.blockState);
+    }, atBottom, stick);
+    assistant.scheduleThinking = makeFrameScheduler(() => {
+      if (!assistant.settled)
+        assistant.thinkBody.textContent = assistant.thinking;
     }, atBottom, stick);
     stick(true);
   }
@@ -3670,6 +3676,7 @@ function createChatController() {
   function finishStreaming(a) {
     if (!a)
       return;
+    a.settled = true;
     const previewIdx = captureCanvasViewStates(a.textNode);
     a.textNode.innerHTML = mdToHtml(a.text);
     restoreCanvasViewStates(a.textNode, previewIdx);

package/src/web/src/chat.ts

@@ -142,6 +142,7 @@ interface AssistantMsgState {
   t0: number;
   tFirst: number;
   tokens: number;
+  settled: boolean;
   blockState: BlockState;
   scheduleText: () => void;
   scheduleThinking: () => void;
@@ -373,7 +374,7 @@ export function createChatController() {
     const meta = el("div", "meta", m);
     curAssistant = {
       m, bubble, thinkBox, thinkBody, textNode, cursor, meta, tools: new Map(),
-      text: "", thinking: "", t0: performance.now(), tFirst: 0, tokens: 0,
+      text: "", thinking: "", t0: performance.now(), tFirst: 0, tokens: 0, settled: false,
       // Block-memoization state (see renderBlocksIncremental) for the main
       // markdown text stream, plus a shared rAF scheduler so a burst of
       // deltas within one frame (e.g. spec-decode emitting several tokens
@@ -390,11 +391,16 @@ export function createChatController() {
     // Consumed — a later turn with no RAG'd attachments must not inherit
     // this turn's citation set (see the pendingCitations declaration above).
     pendingCitations = [];
-    curAssistant.scheduleText = makeFrameScheduler(
-      () => renderBlocksIncremental(curAssistant!.textNode, curAssistant!.text, curAssistant!.blockState),
+    const assistant = curAssistant;
+    assistant.scheduleText = makeFrameScheduler(
+      () => {
+        if (!assistant.settled) renderBlocksIncremental(assistant.textNode, assistant.text, assistant.blockState);
+      },
       atBottom, stick);
-    curAssistant.scheduleThinking = makeFrameScheduler(
-      () => { curAssistant!.thinkBody.textContent = curAssistant!.thinking; },
+    assistant.scheduleThinking = makeFrameScheduler(
+      () => {
+        if (!assistant.settled) assistant.thinkBody.textContent = assistant.thinking;
+      },
       atBottom, stick);
     stick(true);
   }
@@ -612,6 +618,7 @@ export function createChatController() {
    *  directly on the full string. */
   function finishStreaming(a: AssistantMsgState | null): void {
     if (!a) return;
+    a.settled = true;
     // Preserve any Canvas Preview toggle the user made mid-stream (Principle
     // 9: this full non-memoized re-render must not throw away user
     // interaction state) — mdCodeBlock always emits qualifying fences back

package/src/weights.ts

@@ -82,30 +82,7 @@ export class Weights {
     const shards = await (options.openShards ?? ShardedSafetensors.open)(modelDir);
     const self = new Weights(shards, native);
     try {
-      for (const [file, sf] of self.shards.files) {
-        // out-param slots read back via read.u64, not [0] (DFG stale-read
-        // bug — see outArray in mlx/ffi.ts). Cold path, but the rule is
-        // uniform: native wrote it, read.* reads it.
-        const arrMap = new BigUint64Array([native.newArrayMap()]);
-        const arrMapPtr = ptr(arrMap);
-        let arrayMapTransferred = false;
-        try {
-          const metaMap = new BigUint64Array([native.newStringMap()]);
-          const metaMapPtr = ptr(metaMap);
-          try {
-            const status = native.loadSafetensors(arrMap, metaMap, sf.path);
-            if (status !== 0)
-              throw new Error(`mlx_load_safetensors(${sf.path}) failed`);
-            self.#maps.set(file, read.u64(arrMapPtr, 0));
-            arrayMapTransferred = true;
-          } finally {
-            native.freeStringMap(read.u64(metaMapPtr, 0));
-          }
-        } finally {
-          if (!arrayMapTransferred)
-            native.freeArrayMap(read.u64(arrMapPtr, 0));
-        }
-      }
+      for (const [file] of self.shards.files) self.#openShard(file);
       return self;
     } catch (error) {
       // A failure on shard N must release every array map transferred by
@@ -115,6 +92,64 @@ export class Weights {
     }
   }
 
+  /** Load one shard's native string→array map (lazy re-open path too). */
+  #openShard(file: string): void {
+    const sf = this.shards.files.get(file);
+    if (!sf) throw new Error(`no shard file ${file}`);
+    // out-param slots read back via read.u64, not [0] (DFG stale-read
+    // bug — see outArray in mlx/ffi.ts). Cold path, but the rule is
+    // uniform: native wrote it, read.* reads it.
+    const arrMap = new BigUint64Array([this.#native.newArrayMap()]);
+    const arrMapPtr = ptr(arrMap);
+    let arrayMapTransferred = false;
+    try {
+      const metaMap = new BigUint64Array([this.#native.newStringMap()]);
+      const metaMapPtr = ptr(metaMap);
+      try {
+        const status = this.#native.loadSafetensors(arrMap, metaMap, sf.path);
+        if (status !== 0)
+          throw new Error(`mlx_load_safetensors(${sf.path}) failed`);
+        this.#maps.set(file, read.u64(arrMapPtr, 0));
+        arrayMapTransferred = true;
+      } finally {
+        this.#native.freeStringMap(read.u64(metaMapPtr, 0));
+      }
+    } finally {
+      if (!arrayMapTransferred)
+        this.#native.freeArrayMap(read.u64(arrMapPtr, 0));
+    }
+  }
+
+  /** Free one shard's native map AND every cached tensor from it, releasing
+   *  all bytes materialized while consuming that shard. The shard transparently
+   *  re-opens on the next tensor() touching it (header parse only — lazy).
+   *  Required by streaming whole-model transforms at 27B scale: the native map
+   *  retains every evaluated source array, so per-tensor dispose alone frees
+   *  nothing — the whole model accumulates in RAM (2026-08-18 fold OOM). */
+  releaseShard(file: string): void {
+    const sf = this.shards.files.get(file);
+    if (!sf) return;
+    for (const [name, arr] of this.#arrays) {
+      if (this.shards.tensorToFile.get(name) === sf) {
+        arr.dispose();
+        this.#arrays.delete(name);
+      }
+    }
+    const map = this.#maps.get(file);
+    if (map !== undefined) {
+      this.#native.freeArrayMap(map);
+      this.#maps.delete(file);
+    }
+  }
+
+  /** The shard filename holding `name` (for releaseShard scheduling). */
+  fileOf(name: string): string | undefined {
+    const sf = this.shards.tensorToFile.get(name);
+    if (!sf) return undefined;
+    for (const [file, cand] of this.shards.files) if (cand === sf) return file;
+    return undefined;
+  }
+
   get tensorNames(): string[] {
     return this.shards.tensorNames;
   }
@@ -133,8 +168,15 @@ export class Weights {
     if (!arr) {
       const sf = this.shards.tensorToFile.get(name);
       if (!sf) throw new Error(`no tensor named ${name}`);
-      const mapHandle = [...this.#maps.entries()]
-        .find(([file]) => this.shards.files.get(file) === sf)![1];
+      let entry = [...this.#maps.entries()]
+        .find(([file]) => this.shards.files.get(file) === sf);
+      if (!entry) {
+        // Shard was releaseShard()'d — transparently re-open it.
+        const file = this.fileOf(name)!;
+        this.#openShard(file);
+        entry = [file, this.#maps.get(file)!];
+      }
+      const mapHandle = entry[1];
       const slot = new BigUint64Array([C.mlx_array_new()]);
       const slotPtr = ptr(slot);
       if (C.mlx_map_string_to_array_get(slotPtr, mapHandle, ptr(cstr(name))) !== 0)
@@ -145,6 +187,19 @@ export class Weights {
     return arr;
   }
 
+  /** Dispose + evict one cached tensor so its materialized bytes are freed
+   *  mid-walk. Streaming whole-model transforms (fold/quantize at 27B scale)
+   *  MUST release each source tensor after its consumer evaluates, or the
+   *  cache accumulates the entire model in RAM. Safe to call for names never
+   *  accessed; the next tensor(name) call re-opens it from the map. */
+  release(name: string): void {
+    const arr = this.#arrays.get(name);
+    if (arr) {
+      arr.dispose();
+      this.#arrays.delete(name);
+    }
+  }
+
   dispose(): void {
     for (const arr of this.#arrays.values()) arr.dispose();
     this.#arrays.clear();