final

# Conflicts:
#	mlx_tinygrad_interop/stress_interop.py

.envrc

@@ -1 +1,8 @@
 use flake
+
+# creates .venv if doesn't exist and loads its environment
+export VIRTUAL_ENV=".venv"
+if ! [ -d "./$VIRTUAL_ENV" ]; then
+  uv venv
+fi
+layout python

.github/workflows/build-app.yml

@@ -159,7 +159,7 @@ jobs:
           fi
 
       - name: Install Homebrew packages
-        run: brew install just awscli macmon
+        run: brew install just awscli
 
       - name: Install UV
         uses: astral-sh/setup-uv@v6
@@ -239,10 +239,92 @@ jobs:
           # Export keychain path for other steps
           echo "BUILD_KEYCHAIN_PATH=$KEYCHAIN_PATH" >> $GITHUB_ENV
 
+      # ============================================================
+      # Pre-flight credential / profile validation
+      # Runs BEFORE the ~16 min build so auth/expiry failures surface in <1 min.
+      # ============================================================
+
+      - name: Validate Apple notarization credentials
+        env:
+          APPLE_NOTARIZATION_USERNAME: ${{ secrets.APPLE_NOTARIZATION_USERNAME }}
+          APPLE_NOTARIZATION_PASSWORD: ${{ secrets.APPLE_NOTARIZATION_PASSWORD }}
+          APPLE_NOTARIZATION_TEAM: ${{ secrets.APPLE_NOTARIZATION_TEAM }}
+        run: |
+          # All-or-nothing: either all three creds are set, or none are.
+          CRED_COUNT=0
+          for v in "$APPLE_NOTARIZATION_USERNAME" "$APPLE_NOTARIZATION_PASSWORD" "$APPLE_NOTARIZATION_TEAM"; do
+            [[ -n "$v" ]] && CRED_COUNT=$((CRED_COUNT + 1))
+          done
+          if [[ "$CRED_COUNT" -eq 0 ]]; then
+            echo "No notarization credentials configured — skipping notarization for this build."
+            exit 0
+          fi
+          if [[ "$CRED_COUNT" -ne 3 ]]; then
+            echo "ERROR: partial notarization credentials set ($CRED_COUNT/3). Aborting before build."
+            exit 1
+          fi
+          # Cheap, ~5s, auth-only call. Fails instantly with a clear message if
+          # the app-specific password is stale, wrong team-id, etc.
+          echo "Verifying Apple notarization credentials via notarytool history..."
+          if ! xcrun notarytool history \
+            --apple-id "$APPLE_NOTARIZATION_USERNAME" \
+            --password "$APPLE_NOTARIZATION_PASSWORD" \
+            --team-id "$APPLE_NOTARIZATION_TEAM" >/dev/null; then
+            echo "ERROR: notarytool rejected the provided credentials. Fix before rerunning."
+            echo "Common causes: app-specific password expired/revoked, wrong team-id,"
+            echo "Apple ID not on the team, or 2FA not configured for this Apple ID."
+            exit 1
+          fi
+          echo "Apple notarization credentials OK."
+
+      - name: Validate provisioning profile expiry
+        run: |
+          PROFILE="$HOME/Library/Developer/Xcode/UserData/Provisioning Profiles/EXO.provisionprofile"
+          if [[ ! -f "$PROFILE" ]]; then
+            echo "ERROR: provisioning profile not found at $PROFILE"
+            exit 1
+          fi
+          EXPIRY=$(security cms -D -i "$PROFILE" | plutil -extract ExpirationDate raw -o - - 2>/dev/null || true)
+          if [[ -z "$EXPIRY" ]]; then
+            echo "WARNING: could not read ExpirationDate from provisioning profile; skipping expiry check."
+            exit 0
+          fi
+          # Try a couple of known plutil date formats. If none parse, skip the check rather
+          # than risk a false-positive "expired" block on a format we didn't anticipate.
+          EXPIRY_EPOCH=""
+          for fmt in "%Y-%m-%dT%H:%M:%SZ" "%Y-%m-%d %H:%M:%S %z" "%Y-%m-%d %H:%M:%S +0000"; do
+            if parsed=$(date -j -f "$fmt" "$EXPIRY" +%s 2>/dev/null); then
+              EXPIRY_EPOCH="$parsed"
+              break
+            fi
+          done
+          if [[ -z "$EXPIRY_EPOCH" ]]; then
+            echo "WARNING: could not parse ExpirationDate '$EXPIRY'; skipping expiry check."
+            exit 0
+          fi
+          NOW_EPOCH=$(date +%s)
+          if [[ "$EXPIRY_EPOCH" -le "$NOW_EPOCH" ]]; then
+            echo "ERROR: provisioning profile expired on $EXPIRY. Regenerate it before rerunning."
+            exit 1
+          fi
+          DAYS_LEFT=$(( (EXPIRY_EPOCH - NOW_EPOCH) / 86400 ))
+          echo "Provisioning profile valid until $EXPIRY ($DAYS_LEFT days remaining)."
+          if [[ "$DAYS_LEFT" -lt 14 ]]; then
+            echo "WARNING: profile expires in under 14 days — regenerate soon."
+          fi
+
       # ============================================================
       # Build the bundle
       # ============================================================
 
+      - name: Add pinned macmon to PATH
+        run: |
+          MACMON_DIR=$(nix develop --command sh -c 'dirname $(which macmon)')
+          echo "Using macmon from: $MACMON_DIR"
+          echo "$MACMON_DIR" >> $GITHUB_PATH
+          # Remove any Homebrew macmon so PyInstaller can't accidentally pick it up
+          brew uninstall macmon 2>/dev/null || true
+
       - name: Build PyInstaller bundle
         run: uv run pyinstaller packaging/pyinstaller/exo.spec
 
@@ -298,11 +380,41 @@ jobs:
           APPLE_NOTARIZATION_PASSWORD: ${{ secrets.APPLE_NOTARIZATION_PASSWORD }}
           APPLE_NOTARIZATION_TEAM: ${{ secrets.APPLE_NOTARIZATION_TEAM }}
         run: |
+          set -o pipefail
           cd output
           security unlock-keychain -p "$MACOS_CERTIFICATE_PASSWORD" "$BUILD_KEYCHAIN_PATH"
           SIGNING_IDENTITY=$(security find-identity -v -p codesigning "$BUILD_KEYCHAIN_PATH" | awk -F '"' '{print $2}')
+
+          # Fail fast if notarization creds are partial. All-or-nothing.
+          CRED_COUNT=0
+          for v in "$APPLE_NOTARIZATION_USERNAME" "$APPLE_NOTARIZATION_PASSWORD" "$APPLE_NOTARIZATION_TEAM"; do
+            [[ -n "$v" ]] && CRED_COUNT=$((CRED_COUNT + 1))
+          done
+          if [[ "$CRED_COUNT" -ne 0 && "$CRED_COUNT" -ne 3 ]]; then
+            echo "ERROR: partial Apple notarization credentials set ($CRED_COUNT/3). Aborting."
+            exit 1
+          fi
+
           /usr/bin/codesign --deep --force --timestamp --options runtime \
             --sign "$SIGNING_IDENTITY" EXO.app
+
+          # Pre-flight: verify the signed app BEFORE building DMG and submitting to Apple.
+          # If this fails, notarization will fail too — cheap way to fail in seconds, not 15 minutes.
+          echo "===== codesign --verify EXO.app ====="
+          if ! /usr/bin/codesign --verify --deep --strict --verbose=2 EXO.app; then
+            echo "ERROR: EXO.app failed codesign verification. Dumping signing status of every executable:"
+            find EXO.app -type f \( -perm -111 -o -name "*.dylib" -o -name "*.so" -o -name "*.framework" \) -print0 |
+              while IFS= read -r -d '' f; do
+                printf -- '--- %s\n' "$f"
+                /usr/bin/codesign -dv --verbose=2 "$f" 2>&1 | sed 's/^/    /' || true
+              done
+            exit 1
+          fi
+
+          # Gatekeeper assessment. A failure here strongly predicts notarization rejection.
+          echo "===== spctl assessment (predicts notarization outcome) ====="
+          /usr/bin/spctl -a -vvv -t install EXO.app || echo "WARNING: spctl assessment failed — notarization is likely to fail too."
+
           mkdir -p dmg-root
           cp -R EXO.app dmg-root/
           ln -s /Applications dmg-root/Applications
@@ -310,12 +422,22 @@ jobs:
           hdiutil create -volname "EXO" -srcfolder dmg-root -ov -format UDZO "$DMG_NAME"
           /usr/bin/codesign --force --timestamp --options runtime \
             --sign "$SIGNING_IDENTITY" "$DMG_NAME"
+
+          echo "===== codesign --verify DMG ====="
+          if ! /usr/bin/codesign --verify --verbose=2 "$DMG_NAME"; then
+            echo "ERROR: DMG failed codesign verification."
+            exit 1
+          fi
+
           if [[ -n "$APPLE_NOTARIZATION_USERNAME" ]]; then
+            echo "===== notarytool submit ====="
+            # `|| true` so set -e doesn't abort before we can echo output / fetch the log.
+            # We rely on the parsed STATUS below to decide pass/fail.
             SUBMISSION_OUTPUT=$(xcrun notarytool submit "$DMG_NAME" \
               --apple-id "$APPLE_NOTARIZATION_USERNAME" \
               --password "$APPLE_NOTARIZATION_PASSWORD" \
               --team-id "$APPLE_NOTARIZATION_TEAM" \
-              --wait --timeout 15m 2>&1)
+              --wait --timeout 15m 2>&1) || true
             echo "$SUBMISSION_OUTPUT"
 
             SUBMISSION_ID=$(echo "$SUBMISSION_OUTPUT" | awk 'tolower($1)=="id:" && $2 ~ /^[0-9a-fA-F-]+$/ {print $2; exit}')

.github/workflows/pipeline.yml

@@ -91,9 +91,6 @@ jobs:
             nix build .#metal-toolchain
           fi
 
-          # Build mlx (depends on metal-toolchain)
-          nix build .#mlx
-
       - name: Build all Nix outputs
         run: |
           nix flake show --json | jq -r '

.gitignore

@@ -38,3 +38,5 @@ bench/**/*.json
 
 # tmp
 tmp/models
+/build/exo
+/.claude/skills

.idea/exo-v2.iml

@@ -1,31 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<module type="EMPTY_MODULE" version="4">
-  <component name="FacetManager">
-    <facet type="Python" name="Python facet">
-      <configuration sdkName="Python 3.13 virtualenv at ~/Desktop/exo/.venv" />
-    </facet>
-  </component>
-  <component name="Go" enabled="true" />
-  <component name="NewModuleRootManager">
-    <content url="file://$MODULE_DIR$">
-      <sourceFolder url="file://$MODULE_DIR$/scripts/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/exo_pyo3_bindings/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/exo_pyo3_bindings/tests" isTestSource="true" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/util/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/networking/examples" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/networking/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/networking/tests" isTestSource="true" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/system_custodian/src" isTestSource="false" />
-      <excludeFolder url="file://$MODULE_DIR$/.venv" />
-      <excludeFolder url="file://$MODULE_DIR$/.direnv" />
-      <excludeFolder url="file://$MODULE_DIR$/build" />
-      <excludeFolder url="file://$MODULE_DIR$/dist" />
-      <excludeFolder url="file://$MODULE_DIR$/.go_cache" />
-      <excludeFolder url="file://$MODULE_DIR$/rust/target" />
-    </content>
-    <orderEntry type="jdk" jdkName="Python 3.13 (exo)" jdkType="Python SDK" />
-    <orderEntry type="sourceFolder" forTests="false" />
-    <orderEntry type="library" name="Python 3.13 virtualenv at ~/Desktop/exo/.venv interpreter library" level="application" />
-  </component>
-</module>

.idea/externalDependencies.xml

@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
   <component name="ExternalDependencies">
-    <plugin id="systems.fehn.intellijdirenv" />
+    <plugin id="al.aoli.intellijdirenv" />
   </component>
 </project>

.idea/inspectionProfiles/Project_Default.xml

@@ -1,14 +0,0 @@
-<component name="InspectionProjectProfileManager">
-  <profile version="1.0">
-    <option name="myName" value="Project Default" />
-    <inspection_tool class="PyCompatibilityInspection" enabled="true" level="WARNING" enabled_by_default="true">
-      <option name="ourVersions">
-        <value>
-          <list size="1">
-            <item index="0" class="java.lang.String" itemvalue="3.14" />
-          </list>
-        </value>
-      </option>
-    </inspection_tool>
-  </profile>
-</component>

.idea/misc.xml

@@ -4,7 +4,4 @@
     <option name="sdkName" value="Python 3.13 (exo)" />
   </component>
   <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.13 (exo)" project-jdk-type="Python SDK" />
-  <component name="PythonCompatibilityInspectionAdvertiser">
-    <option name="version" value="3" />
-  </component>
 </project>

.mlx_typings/mlx/core/__init__.pyi

@@ -1767,12 +1767,12 @@ def clip(
         array: The clipped array.
     """
 
-def compile(
-    fun: Callable,
+def compile[F: Callable[..., object]](
+    fun: F,
     inputs: object | None = ...,
     outputs: object | None = ...,
     shapeless: bool = ...,
-) -> Callable:
+) -> F:
     """
     Returns a compiled function which produces the same output as ``fun``.
 
@@ -2396,7 +2396,7 @@ def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
         array: The angles in degrees.
     """
 
-def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]):
+def depends[T](inputs: T, dependencies: array | Sequence[array]) -> T:
     """
     Insert dependencies between arrays in the graph. The outputs are
     identical to ``inputs`` but with dependencies on ``dependencies``.
@@ -2915,8 +2915,8 @@ def gather_mm(
     a: array,
     b: array,
     /,
-    lhs_indices: array,
-    rhs_indices: array,
+    lhs_indices: array | None = ...,
+    rhs_indices: array | None = ...,
     *,
     sorted_indices: bool = ...,
     stream: Stream | Device | None = ...,
@@ -4707,6 +4707,7 @@ def softmax(
     /,
     axis: int | Sequence[int] | None = ...,
     *,
+    precise: bool = ...,
     stream: Stream | Device | None = ...,
 ) -> array:
     """

.mlx_typings/mlx/nn/__init__.pyi

@@ -1,9 +1,5 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from layers import *
-from utils import *
+from .layers import *
+from .utils import *
 
 from . import init as init
 from . import losses as losses

.mlx_typings/mlx/nn/layers/__init__.pyi

@@ -1,20 +1,16 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from activations import *
-from base import *
-from containers import *
-from convolution import *
-from convolution_transpose import *
-from distributed import *
-from dropout import *
-from embedding import *
-from linear import *
-from normalization import *
-from pooling import *
-from positional_encoding import *
-from quantized import *
-from recurrent import *
-from transformer import *
-from upsample import *
+from .activations import *
+from .base import *
+from .containers import *
+from .convolution import *
+from .convolution_transpose import *
+from .distributed import *
+from .dropout import *
+from .embedding import *
+from .linear import *
+from .normalization import *
+from .pooling import *
+from .positional_encoding import *
+from .quantized import *
+from .recurrent import *
+from .transformer import *
+from .upsample import *

.mlx_typings/mlx/nn/layers/base.pyi

@@ -53,10 +53,14 @@ class Module(dict):
         mx.eval(model.parameters())
     """
 
-    __call__: Callable
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
     def __init__(self) -> None:
         """Should be called by the subclasses of ``Module``."""
 
+    def __getitem__(self, key: str) -> mx.array | Module: ...
+    def get(
+        self, key: str, default: mx.array | Module | None = ...
+    ) -> mx.array | Module | None: ...
     @property
     def training(self):  # -> bool:
         """Boolean indicating if the model is in training mode."""

.mlx_typings/mlx/nn/layers/quantized.pyi

@@ -2,7 +2,7 @@
 This type stub file was generated by pyright.
 """
 
-from typing import Callable, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import mlx.core as mx
 from base import Module
@@ -13,8 +13,10 @@ def quantize(
     bits: int = ...,
     *,
     mode: str = ...,
-    class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = ...,
-):  # -> None:
+    class_predicate: Optional[
+        Callable[[str, Module], Union[bool, dict[str, Any]]]
+    ] = ...,
+) -> None:
     """Quantize the sub-modules of a module according to a predicate.
 
     By default all layers that define a ``to_quantized(group_size, bits)``

.mlx_typings/mlx_lm/_version.pyi

@@ -0,0 +1,5 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+__version__ = ...

.mlx_typings/mlx_lm/generate.pyi

@@ -3,13 +3,12 @@ This type stub file was generated by pyright.
 """
 
 import contextlib
-from dataclasses import dataclass
-from typing import Any, Callable, Generator, List, Optional, Tuple, Union
-
 import mlx.core as mx
 import mlx.nn as nn
+from dataclasses import dataclass
+from collections import deque
+from typing import Any, Callable, Generator, List, Optional, Sequence, Tuple, Union
 from transformers import PreTrainedTokenizer
-
 from .tokenizer_utils import TokenizerWrapper
 
 DEFAULT_PROMPT = ...
@@ -29,8 +28,9 @@ def str2bool(string):  # -> bool:
     ...
 def setup_arg_parser():  # -> ArgumentParser:
     """Set up and return the argument parser."""
+    ...
 
-generation_stream = ...
+generation_stream: mx.Stream
 
 @contextlib.contextmanager
 def wired_limit(
@@ -43,6 +43,7 @@ def wired_limit(
     async eval could be running pass in the streams to synchronize with prior
     to exiting the context manager.
     """
+    ...
 @dataclass
 class GenerationResponse:
     """
@@ -91,7 +92,7 @@ def generate_step(
     kv_bits: Optional[int] = ...,
     kv_group_size: int = ...,
     quantized_kv_start: int = ...,
-    prompt_progress_callback: Optional[Callable[[int], int]] = ...,
+    prompt_progress_callback: Optional[Callable[[int, int], None]] = ...,
     input_embeddings: Optional[mx.array] = ...,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
     """
@@ -117,7 +118,7 @@ def generate_step(
         kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
         quantized_kv_start (int): Step to begin using a quantized KV cache.
            when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_progress_callback (Callable[[int], int]): A call-back which takes the
+        prompt_progress_callback (Callable[[int, int], None]): A call-back which takes the
            prompt tokens processed so far and the total number of prompt tokens.
         input_embeddings (mx.array, optional): Input embeddings to use instead of or in
           conjunction with prompt tokens. Default: ``None``.
@@ -125,6 +126,7 @@ def generate_step(
     Yields:
         Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
     """
+    ...
 
 def speculative_generate_step(
     prompt: mx.array,
@@ -170,6 +172,7 @@ def speculative_generate_step(
         Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
           and a bool indicating if the token was generated by the draft model
     """
+    ...
 
 def stream_generate(
     model: nn.Module,
@@ -177,7 +180,7 @@ def stream_generate(
     prompt: Union[str, mx.array, List[int]],
     max_tokens: int = ...,
     draft_model: Optional[nn.Module] = ...,
-    **kwargs: object,
+    **kwargs: Any,
 ) -> Generator[GenerationResponse, None, None]:
     """
     A generator producing text based on the given prompt from the model.
@@ -199,6 +202,7 @@ def stream_generate(
         GenerationResponse: An instance containing the generated text segment and
             associated metadata. See :class:`GenerationResponse` for details.
     """
+    ...
 
 def generate(
     model: nn.Module,
@@ -219,6 +223,9 @@ def generate(
        kwargs: The remaining options get passed to :func:`stream_generate`.
           See :func:`stream_generate` for more details.
     """
+    ...
+
+def _merge_caches(caches: List[List[Any]]) -> List[Any]: ...
 @dataclass
 class BatchStats:
     """
@@ -242,10 +249,262 @@ class BatchStats:
     generation_time: float = ...
     peak_memory: float = ...
 
+class SequenceStateMachine:
+    """A state machine that uses one Aho-Corasick trie per state to efficiently
+    track state across a generated sequence.
+
+    The transitions are provided as state -> [(sequence, new_state)].
+
+    Example:
+
+        sm = SequenceStateMachine(
+            transitions={
+                "normal": [
+                    (think_start_tokens, "reasoning"),
+                    (tool_start_tokens, "tool"),
+                    (eos, None),
+                ],
+                "reasoning": [
+                    (think_end_tokens, "normal"),
+                    (eos, None),
+                ],
+                "tool": [
+                    (tool_end_tokens, None),
+                    (eos, None)
+                ],
+            },
+            initial="normal"
+        )
+    """
+    def __init__(self, transitions=..., initial=...) -> None: ...
+    def __deepcopy__(self, memo):  # -> SequenceStateMachine:
+        ...
+    def make_state(self):  # -> tuple[str, Any, dict[Any, Any]]:
+        ...
+    @staticmethod
+    def match(state, x):  # -> tuple[tuple[Any, Any | None, Any], Any | None, Any]:
+        ...
+
+class PromptProcessingBatch:
+    """
+    A batch processor for prompt tokens with support for incremental processing.
+
+    This class handles batched prompt processing, managing KV caches and preparing
+    tokens for generation. It supports extending, filtering, and splitting batches.
+    """
+    @dataclass
+    class Response:
+        uid: int
+        progress: tuple
+        end_of_segment: bool
+        end_of_prompt: bool
+        ...
+
+    def __init__(
+        self,
+        model: nn.Module,
+        uids: List[int],
+        caches: List[List[Any]],
+        tokens: Optional[List[List[int]]] = ...,
+        prefill_step_size: int = ...,
+        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+        fallback_sampler: Optional[Callable[[mx.array], mx.array]] = ...,
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ] = ...,
+        state_machines: Optional[List[SequenceStateMachine]] = ...,
+        max_tokens: Optional[List[int]] = ...,
+    ) -> None: ...
+    def __len__(self):  # -> int:
+        ...
+    def extract_cache(self, idx: int) -> List[Any]: ...
+    def extend(self, batch):  # -> None:
+        ...
+    def split(self, indices: List[int]):  # -> Self:
+        ...
+    def filter(self, keep: List[int]):  # -> None:
+        ...
+    def prompt(self, tokens: List[List[int]]):  # -> None:
+        """
+        Process prompt tokens through the model.
+
+        Args:
+            tokens: List of token sequences to process.
+        """
+        ...
+
+    def generate(self, tokens: List[List[int]]):  # -> GenerationBatch:
+        """
+        Transition from prompt processing to generation.
+
+        Args:
+            tokens: Final tokens for each sequence to start generation.
+
+        Returns:
+            A GenerationBatch ready for token generation.
+        """
+        ...
+
+    @classmethod
+    def empty(
+        cls,
+        model: nn.Module,
+        fallback_sampler: Callable[[mx.array], mx.array],
+        prefill_step_size: int = ...,
+    ):  # -> Self:
+        ...
+
+class GenerationBatch:
+    """
+    A batched token generator that manages multiple sequences in parallel.
+
+    This class handles the generation phase after prompt processing, managing
+    KV caches, sampling, and stop sequence detection for multiple sequences.
+    """
+    @dataclass
+    class Response:
+        uid: int
+        token: int
+        logprobs: mx.array
+        finish_reason: Optional[str]
+        current_state: Optional[str]
+        match_sequence: Optional[List[int]]
+        prompt_cache: Optional[List[Any]]
+        all_tokens: Optional[List[int]]
+        ...
+
+    model: nn.Module
+    uids: List[int]
+    prompt_cache: List[Any]
+    tokens: List[List[int]]
+    samplers: Optional[List[Callable[[mx.array], mx.array]]]
+    fallback_sampler: Callable[[mx.array], mx.array]
+    logits_processors: Optional[List[List[Callable[[mx.array, mx.array], mx.array]]]]
+    state_machines: List[SequenceStateMachine]
+    max_tokens: List[int]
+    _current_tokens: Optional[mx.array]
+    _current_logprobs: List[mx.array]
+    _next_tokens: mx.array
+    _next_logprobs: List[mx.array]
+    _token_context: List[mx.array]
+    _num_tokens: List[int]
+
+    def __init__(
+        self,
+        model: nn.Module,
+        uids: List[int],
+        inputs: mx.array,
+        prompt_cache: List[Any],
+        tokens: List[List[int]],
+        samplers: Optional[List[Callable[[mx.array], mx.array]]],
+        fallback_sampler: Callable[[mx.array], mx.array],
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ],
+        state_machines: List[SequenceStateMachine],
+        max_tokens: List[int],
+    ) -> None: ...
+    def __len__(self) -> int: ...
+    def extend(self, batch: GenerationBatch) -> None: ...
+    def extract_cache(self, idx: int) -> List[Any]: ...
+    def filter(self, keep: List[int]) -> None: ...
+    def _step(self) -> Tuple[List[int], List[mx.array]]: ...
+    def next(self) -> List[Response]:
+        """
+        Generate the next batch of tokens.
+
+        Returns:
+            List of Response objects for each sequence in the batch.
+        """
+        ...
+
+    @classmethod
+    def empty(
+        cls, model: nn.Module, fallback_sampler: Callable[[mx.array], mx.array]
+    ):  # -> Self:
+        ...
+
+class BatchGenerator:
+    """
+    A batch generator implements continuous batching.
+
+    This class provides automatic management of prompt processing and generation
+    batches, handling the transition between the two.
+
+    It also allows for segmented prompt processing which guarantees that the
+    generator will stop at these boundaries when processing an input.
+    """
+    def __init__(
+        self,
+        model: nn.Module,
+        max_tokens: int = ...,
+        stop_tokens: Optional[Sequence[Sequence[int]]] = ...,
+        sampler: Optional[Callable[[mx.array], mx.array]] = ...,
+        logits_processors: Optional[
+            List[Callable[[mx.array, mx.array], mx.array]]
+        ] = ...,
+        completion_batch_size: int = ...,
+        prefill_batch_size: int = ...,
+        prefill_step_size: int = ...,
+    ) -> None: ...
+    def close(self) -> None: ...
+    def __del__(self):  # -> None:
+        ...
+    @contextlib.contextmanager
+    def stats(self, stats=...):  # -> Generator[Any | BatchStats, Any, None]:
+        ...
+    _unprocessed_sequences: deque[tuple[Any, ...]]
+    _prompt_batch: PromptProcessingBatch
+    _generation_batch: GenerationBatch
+    _currently_processing: list[Any]
+    _gen_tokens_counter: int
+    _steps_counter: int
+    def _next(
+        self,
+    ) -> tuple[
+        List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
+    ]: ...
+    def insert(
+        self,
+        prompts: List[List[int]],
+        max_tokens: Optional[List[int]] = ...,
+        caches: Optional[List[List[Any]]] = ...,
+        all_tokens: Optional[List[List[int]]] = ...,
+        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ] = ...,
+        state_machines: Optional[List[SequenceStateMachine]] = ...,
+    ) -> List[int]: ...
+    def insert_segments(
+        self,
+        segments: List[List[List[int]]],
+        max_tokens: Optional[List[int]] = ...,
+        caches: Optional[List[List[Any]]] = ...,
+        all_tokens: Optional[List[List[int]]] = ...,
+        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ] = ...,
+        state_machines: Optional[List[SequenceStateMachine]] = ...,
+    ) -> List[int]: ...
+    def extract_cache(self, uids: List[int]) -> dict[int, Any]: ...
+    def remove(
+        self, uids: List[int], return_prompt_caches: bool = ...
+    ) -> dict[int, Any]: ...
+    @property
+    def prompt_cache_nbytes(self) -> int: ...
+    def next(
+        self,
+    ) -> tuple[
+        List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
+    ]: ...
+    def next_generated(self) -> List[GenerationBatch.Response]: ...
+
 @dataclass
 class BatchResponse:
     """
-    An data object to hold a batch generation response.
+    A data object to hold a batch generation response.
 
     Args:
         texts: (List[str]): The generated text for each prompt.
@@ -255,93 +514,17 @@ class BatchResponse:
     texts: List[str]
     stats: BatchStats
     caches: Optional[List[List[Any]]]
-
-def _left_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
-def _right_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
-def _make_cache(
-    model: Any, left_padding: Any, max_kv_size: Optional[int]
-) -> List[Any]: ...
-def _merge_caches(caches: Any) -> List[Any]: ...
-@dataclass
-class Batch:
-    uids: List[int]
-    y: mx.array
-    logprobs: mx.array
-    max_tokens: List[int]
-    num_tokens: List[int]
-    cache: List[Any]
-    samplers: List[Any]
-    logits_processors: List[Any]
-    tokens: List[mx.array]
-    def __len__(self) -> int: ...
-    def filter(self, keep_idx: List[int]) -> None: ...
-    def extend(self, other: "Batch") -> None: ...
-    def extract_cache(self, idx: int) -> List[Any]: ...
-
-class BatchGenerator:
-    model: Any
-    max_kv_size: Optional[int]
-    prefill_step_size: int
-    unprocessed_prompts: List[Any]
-    active_batch: Optional[Batch]
-    prompt_progress_callback: Callable[[List[Tuple[int, int, int]]], None]
-    _stats: BatchStats
-
-    @dataclass
-    class Response:
-        uid: int
-        token: int
-        logprobs: mx.array
-        finish_reason: Optional[str]
-        prompt_cache: Any
-
-    def __init__(
-        self,
-        model: Any,
-        max_tokens: int = ...,
-        stop_tokens: Optional[set[int]] = ...,
-        sampler: Optional[Callable[[mx.array], mx.array]] = ...,
-        logits_processors: Optional[
-            List[Callable[[mx.array, mx.array], mx.array]]
-        ] = ...,
-        completion_batch_size: int = ...,
-        prefill_batch_size: int = ...,
-        prefill_step_size: int = ...,
-        prompt_progress_callback: Optional[
-            Callable[[List[Tuple[int, int, int]]], None]
-        ] = ...,
-        max_kv_size: Optional[int] = ...,
-    ) -> None: ...
-    def close(self) -> None: ...
-    def insert(
-        self,
-        prompts: Any,
-        max_tokens: Union[List[int], int, None] = ...,
-        caches: Any = ...,
-        samplers: Optional[List[Any]] = ...,
-        logits_processors: Optional[List[Any]] = ...,
-    ) -> List[int]: ...
-    def remove(
-        self, uids: List[int], return_prompt_caches: bool = ...
-    ) -> Optional[dict[int, List[Any]]]: ...
-    def stats(self) -> BatchStats: ...
-    def next(self) -> List[Response]: ...
-    def _process_prompts(self, prompts: List[Any]) -> Batch: ...
-    def _step(
-        self,
-        input_tokens: mx.array,
-        prompt_cache: List[Any],
-        samplers: Optional[List[Any]],
-        logits_processors: Optional[List[Any]],
-        tokens: List[mx.array],
-    ) -> Tuple[mx.array, List[mx.array]]: ...
+    ...
 
 def batch_generate(
     model,
     tokenizer,
-    prompts: List[int],
+    prompts: List[List[int]],
+    prompt_caches: Optional[List[List[Any]]] = ...,
     max_tokens: Union[int, List[int]] = ...,
     verbose: bool = ...,
+    return_prompt_caches: bool = ...,
+    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = ...,
     **kwargs,
 ) -> BatchResponse:
     """
@@ -350,14 +533,22 @@ def batch_generate(
     Args:
        model (nn.Module): The language model.
        tokenizer (PreTrainedTokenizer): The tokenizer.
-       prompt (List[List[int]]): The input prompts.
+       prompts (List[List[int]]): The input prompts.
+       prompt_caches (List[List[Any]], optional): Pre-computed prompt-caches
+          for each input prompt. Note, unlike ``generate_step``, the caches
+          won't be updated in-place.
        verbose (bool): If ``True``, print tokens and timing information.
           Default: ``False``.
        max_tokens (Union[int, List[int]): Maximum number of output tokens. This
           can be per prompt if a list is provided.
+       return_prompt_caches (bool): Return the prompt caches in the batch
+          responses. Default: ``False``.
+       logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
+          A list of functions that take tokens and logits and return the processed logits. Default: ``None``.
        kwargs: The remaining options get passed to :obj:`BatchGenerator`.
           See :obj:`BatchGenerator` for more details.
     """
+    ...
 
 def main():  # -> None:
     ...

.mlx_typings/mlx_lm/models/activations.pyi

@@ -0,0 +1,19 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+from functools import partial
+
+@partial(mx.compile, shapeless=True)
+def swiglu(gate, x): ...
+@partial(mx.compile, shapeless=True)
+def xielu(x, alpha_p, alpha_n, beta, eps):  # -> array:
+    ...
+
+class XieLU(nn.Module):
+    def __init__(
+        self, alpha_p_init=..., alpha_n_init=..., beta=..., eps=...
+    ) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...

.mlx_typings/mlx_lm/models/base.pyi

@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
 """
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional
 
 import mlx.core as mx
 
@@ -37,10 +37,10 @@ def quantized_scaled_dot_product_attention(
     bits: int = ...,
 ) -> mx.array: ...
 def scaled_dot_product_attention(
-    queries,
-    keys,
-    values,
-    cache,
+    queries: mx.array,
+    keys: mx.array,
+    values: mx.array,
+    cache: Optional[Any],
     scale: float,
     mask: Optional[mx.array],
     sinks: Optional[mx.array] = ...,

.mlx_typings/mlx_lm/models/cache.pyi

@@ -88,8 +88,8 @@ def create_attention_mask(
 ) -> array | Literal["causal"] | None: ...
 
 class _BaseCache(Cache):
-    keys: mx.array
-    values: mx.array
+    keys: mx.array | None
+    values: mx.array | None
     offset: int
     @property
     def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@@ -268,29 +268,14 @@ class CacheList(_BaseCache):
         """
 
 class BatchKVCache(_BaseCache):
-    step = ...
-    def __init__(self, left_padding: List[int]) -> None:
-        """
-        The BatchKV cache expects inputs to be left-padded.
-
-        E.g. the following prompts:
-
-            [1, 3, 5]
-            [7]
-            [2, 6, 8, 9]
-
-        Should be padded like so:
-
-            [0, 1, 3, 5]
-            [0, 0, 0, 7]
-            [2, 6, 8, 9]
-
-        And ``left_padding`` specifies the amount of padding for each.
-        In this case, ``left_padding = [1, 3, 0]``.
-        """
-
-    def update_and_fetch(self, keys, values):  # -> tuple[array | Any, array | Any]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    _idx: int
+    def __init__(self, left_padding: List[int]) -> None: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
     @property
     def state(
         self,
@@ -316,12 +301,21 @@ class BatchKVCache(_BaseCache):
         """
 
 class BatchRotatingKVCache(_BaseCache):
-    step = ...
-    def __init__(self, max_size, left_padding: List[int]) -> None: ...
-    def update_and_fetch(
-        self, keys, values
-    ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    max_size: int
+    _idx: int
+    _offset: int
+    rotated: bool
+    _lengths: array | None
+    def __init__(self, max_size: int, left_padding: List[int]) -> None: ...
+    def _trim(self, trim_size: int, v: array, append: array | None = ...) -> array: ...
+    def _update_in_place(self, keys: array, values: array) -> tuple[array, array]: ...
+    def _update_concat(self, keys: array, values: array) -> tuple[array, array]: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
     @property
     def state(
         self,

.mlx_typings/mlx_lm/models/gated_delta.pyi

@@ -0,0 +1,35 @@
+from typing import Optional
+
+import mlx.core as mx
+
+def compute_g(A_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array: ...
+def gated_delta_update(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    a: mx.array,
+    b: mx.array,
+    A_log: mx.array,
+    dt_bias: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+    use_kernel: bool = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_ops(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_kernel(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...

.mlx_typings/mlx_lm/models/gemma4.pyi

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from . import gemma4_text
+from .base import BaseModelArgs
+from .cache import KVCache, RotatingKVCache
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: Optional[dict[str, Any]]
+    vocab_size: int
+
+    def __post_init__(self) -> None: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    language_model: gemma4_text.Model
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[gemma4_text.DecoderLayer]: ...
+    @property
+    def quant_predicate(self) -> Any: ...
+    def make_cache(self) -> list[KVCache | RotatingKVCache]: ...

.mlx_typings/mlx_lm/models/gemma4_text.pyi

@@ -0,0 +1,179 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import KVCache, RotatingKVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    head_dim: int
+    global_head_dim: int
+    global_partial_rotary_factor: float
+    rms_norm_eps: float
+    vocab_size: int
+    vocab_size_per_layer_input: int
+    num_key_value_heads: int
+    num_global_key_value_heads: Optional[int]
+    num_kv_shared_layers: int
+    pad_token_id: int
+    hidden_size_per_layer_input: int
+    rope_traditional: bool
+    partial_rotary_factor: float
+    rope_parameters: Optional[Dict[str, Any]]
+    sliding_window: int
+    sliding_window_pattern: int
+    max_position_embeddings: int
+    attention_k_eq_v: bool
+    final_logit_softcapping: float
+    use_double_wide_mlp: bool
+    enable_moe_block: bool
+    num_experts: Optional[int]
+    top_k_experts: Optional[int]
+    moe_intermediate_size: Optional[int]
+    layer_types: Optional[List[str]]
+    tie_word_embeddings: bool
+
+    def __post_init__(self) -> None: ...
+
+class MLP(nn.Module):
+    gate_proj: nn.Linear
+    down_proj: nn.Linear
+    up_proj: nn.Linear
+
+    def __init__(self, config: ModelArgs, layer_idx: int = 0) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class Router(nn.Module):
+    proj: nn.Linear
+    scale: mx.array
+    per_expert_scale: mx.array
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
+class Experts(nn.Module):
+    switch_glu: SwitchGLU
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self, x: mx.array, top_k_indices: mx.array, top_k_weights: mx.array
+    ) -> mx.array: ...
+
+class Attention(nn.Module):
+    layer_idx: int
+    layer_type: str
+    is_sliding: bool
+    head_dim: int
+    n_heads: int
+    n_kv_heads: int
+    use_k_eq_v: bool
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    q_norm: nn.Module
+    k_norm: nn.Module
+    v_norm: nn.Module
+    rope: nn.Module
+
+    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+class DecoderLayer(nn.Module):
+    layer_idx: int
+    layer_type: str
+    self_attn: Attention
+    mlp: MLP
+    enable_moe: bool
+    router: Router
+    experts: Experts
+    input_layernorm: nn.Module
+    post_attention_layernorm: nn.Module
+    pre_feedforward_layernorm: nn.Module
+    post_feedforward_layernorm: nn.Module
+    post_feedforward_layernorm_1: nn.Module
+    post_feedforward_layernorm_2: nn.Module
+    pre_feedforward_layernorm_2: nn.Module
+    hidden_size_per_layer_input: int
+    per_layer_input_gate: Optional[nn.Linear]
+    per_layer_projection: Optional[nn.Linear]
+    post_per_layer_input_norm: Optional[nn.Module]
+    layer_scalar: mx.array
+
+    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = ...,
+        cache: Optional[Any] = ...,
+        per_layer_input: Optional[mx.array] = ...,
+        shared_kv: Optional[tuple[mx.array, mx.array]] = ...,
+        offset: Optional[mx.array] = ...,
+    ) -> tuple[mx.array, tuple[mx.array, mx.array], mx.array]: ...
+
+class Gemma4TextModel(nn.Module):
+    config: ModelArgs
+    vocab_size: int
+    window_size: int
+    sliding_window_pattern: int
+    num_hidden_layers: int
+    embed_tokens: nn.Embedding
+    embed_scale: float
+    layers: list[DecoderLayer]
+    norm: nn.Module
+    hidden_size_per_layer_input: int
+    embed_tokens_per_layer: Optional[nn.Embedding]
+    per_layer_model_projection: Optional[nn.Linear]
+    per_layer_projection_norm: Optional[nn.Module]
+    previous_kvs: list[int]
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: Optional[mx.array] = ...,
+        cache: Optional[list[Any]] = ...,
+        input_embeddings: Optional[mx.array] = ...,
+        per_layer_inputs: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _get_per_layer_inputs(
+        self,
+        input_ids: Optional[mx.array],
+        input_embeddings: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _project_per_layer_inputs(
+        self,
+        input_embeddings: mx.array,
+        per_layer_inputs: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _make_masks(self, h: mx.array, cache: list[Any]) -> list[Any]: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    model: Gemma4TextModel
+    final_logit_softcapping: float
+    tie_word_embeddings: bool
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[DecoderLayer]: ...
+    @property
+    def head_dim(self) -> int: ...
+    @property
+    def n_kv_heads(self) -> int: ...
+    @property
+    def quant_predicate(self) -> Any: ...
+    def make_cache(self) -> list[KVCache | RotatingKVCache]: ...

.mlx_typings/mlx_lm/models/gpt_oss.pyi

@@ -0,0 +1,103 @@
+"""Type stubs for mlx_lm.models.gpt_oss"""
+
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import KVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    vocab_size: int
+    rms_norm_eps: float
+    sliding_window: int
+    layer_types: Optional[List[str]]
+
+def mlx_topk(a: mx.array, k: int, axis: int = -1) -> tuple[mx.array, mx.array]: ...
+
+class AttentionBlock(nn.Module):
+    head_dim: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_key_value_groups: int
+    sinks: mx.array
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    sm_scale: float
+    rope: nn.Module
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class TransformerBlock(nn.Module):
+    self_attn: AttentionBlock
+    mlp: MLPBlock
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MLPBlock(nn.Module):
+    hidden_size: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    experts: SwitchGLU
+    router: nn.Linear
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class GptOssMoeModel(nn.Module):
+    embed_tokens: nn.Embedding
+    norm: nn.RMSNorm
+    layer_types: List[str]
+    layers: list[TransformerBlock]
+    window_size: int
+    swa_idx: int
+    ga_idx: int
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: GptOssMoeModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[nn.Module]: ...
+    def make_cache(self) -> list[KVCache]: ...

.mlx_typings/mlx_lm/models/minimax.pyi

@@ -0,0 +1,94 @@
+"""Type stubs for mlx_lm.models.minimax"""
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    max_position_embeddings: int
+
+class MiniMaxAttention(nn.Module):
+    num_heads: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    q_norm: nn.Module
+    k_norm: nn.Module
+    rope: nn.Module
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxSparseMoeBlock(nn.Module):
+    num_experts_per_tok: int
+    gate: nn.Linear
+    switch_mlp: SwitchGLU
+    e_score_correction_bias: mx.array
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class MiniMaxDecoderLayer(nn.Module):
+    self_attn: MiniMaxAttention
+    block_sparse_moe: MiniMaxSparseMoeBlock
+    input_layernorm: nn.RMSNorm
+    post_attention_layernorm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxModel(nn.Module):
+    embed_tokens: nn.Embedding
+    layers: list[MiniMaxDecoderLayer]
+    norm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: MiniMaxModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[MiniMaxDecoderLayer]: ...

.mlx_typings/mlx_lm/models/nemotron_h.pyi

@@ -92,6 +92,15 @@ class NemotronHAttention(nn.Module):
         cache: Optional[KVCache] = None,
     ) -> mx.array: ...
 
+class MoEGate(nn.Module):
+    config: ModelArgs
+    top_k: int
+    norm_topk_prob: bool
+    weight: mx.array
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
 class NemotronHMLP(nn.Module):
     up_proj: nn.Linear
     down_proj: nn.Linear
@@ -102,9 +111,14 @@ class NemotronHMLP(nn.Module):
     def __call__(self, x: mx.array) -> mx.array: ...
 
 class NemotronHMoE(nn.Module):
+    config: ModelArgs
     num_experts_per_tok: int
+    moe_latent_size: Optional[int]
     switch_mlp: SwitchMLP
+    gate: MoEGate
     shared_experts: NemotronHMLP
+    fc1_latent_proj: nn.Linear
+    fc2_latent_proj: nn.Linear
 
     def __init__(self, config: ModelArgs) -> None: ...
     def __call__(self, x: mx.array) -> mx.array: ...

.mlx_typings/mlx_lm/models/qwen3_next.pyi

@@ -71,6 +71,7 @@ class Qwen3NextAttention(nn.Module):
 class Qwen3NextSparseMoeBlock(nn.Module):
     norm_topk_prob: bool
     num_experts: int
+    num_experts_per_tok: int
     top_k: int
     gate: nn.Linear
     switch_mlp: SwitchGLU

.mlx_typings/mlx_lm/models/rope_utils.pyi

@@ -0,0 +1,51 @@
+from typing import Any, Optional
+
+import mlx.nn as nn
+
+class YarnRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        beta_fast: float = ...,
+        beta_slow: float = ...,
+        mscale: float = ...,
+        mscale_all_dim: float = ...,
+    ) -> None: ...
+
+class Llama3RoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        low_freq_factor: float = ...,
+        high_freq_factor: float = ...,
+    ) -> None: ...
+
+class SuScaledRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        short_factor: Any = ...,
+        long_factor: Any = ...,
+        original_max_position_embeddings: int = ...,
+    ) -> None: ...
+
+def initialize_rope(
+    dims: int,
+    base: float = ...,
+    traditional: bool = ...,
+    scaling_config: Optional[dict[str, Any]] = ...,
+    max_position_embeddings: Optional[int] = ...,
+) -> nn.Module: ...

.mlx_typings/mlx_lm/sample_utils.pyi

@@ -48,6 +48,10 @@ def make_logits_processors(
     logit_bias: Optional[Dict[int, float]] = ...,
     repetition_penalty: Optional[float] = ...,
     repetition_context_size: Optional[int] = ...,
+    presence_penalty: Optional[float] = ...,
+    presence_context_size: Optional[int] = ...,
+    frequency_penalty: Optional[float] = ...,
+    frequency_context_size: Optional[int] = ...,
 ) -> list[Callable[[mx.array, mx.array], mx.array]]:
     """
     Make logits processors for use with ``generate_step``.

.mlx_typings/mlx_lm/tokenizer_utils.pyi

@@ -117,6 +117,8 @@ class TokenizerWrapper:
     think_end: str | None
     think_start_id: int | None
     think_end_id: int | None
+    think_start_tokens: list[int] | None
+    think_end_tokens: list[int] | None
 
     def __init__(
         self,

.mlx_typings/mlx_lm/tuner/dora.pyi

@@ -0,0 +1,50 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+
+class DoRALinear(nn.Module):
+    @staticmethod
+    def from_base(
+        linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> DoRALinear:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedLinear | Linear:
+        ...
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+        bias: bool = ...,
+    ) -> None: ...
+    def set_linear(self, linear):  # -> None:
+        """
+        Set the self.linear layer and recompute self.m.
+        """
+        ...
+
+    def __call__(self, x): ...
+
+class DoRAEmbedding(nn.Module):
+    def from_base(
+        embedding: nn.Embedding, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> DoRAEmbedding:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> Embedding:
+        ...
+    def __init__(
+        self,
+        num_embeddings: int,
+        dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+    ) -> None: ...
+    def set_embedding(self, embedding: nn.Module):  # -> None:
+        ...
+    def __call__(self, x): ...
+    def as_linear(self, x): ...

.mlx_typings/mlx_lm/tuner/lora.pyi

@@ -0,0 +1,66 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+
+class LoRALinear(nn.Module):
+    @staticmethod
+    def from_base(
+        linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> LoRALinear:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedLinear | Linear:
+        ...
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+        bias: bool = ...,
+    ) -> None: ...
+    def __call__(self, x):  # -> array:
+        ...
+
+class LoRASwitchLinear(nn.Module):
+    @staticmethod
+    def from_base(
+        linear: nn.Module, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> LoRASwitchLinear:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedSwitchLinear | SwitchLinear:
+        ...
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        num_experts: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+        bias: bool = ...,
+    ) -> None: ...
+    def __call__(self, x, indices, sorted_indices=...): ...
+
+class LoRAEmbedding(nn.Module):
+    @staticmethod
+    def from_base(
+        embedding: nn.Embedding, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> LoRAEmbedding:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedEmbedding | Embedding:
+        ...
+    def __init__(
+        self,
+        num_embeddings: int,
+        dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+    ) -> None: ...
+    def __call__(self, x):  # -> array:
+        ...
+    def as_linear(self, x):  # -> array:
+        ...

.mlx_typings/mlx_lm/tuner/utils.pyi

@@ -0,0 +1,57 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+from typing import Dict
+
+def build_schedule(schedule_config: Dict):  # -> Any:
+    """
+    Build a learning rate schedule from the given config.
+    """
+    ...
+
+def linear_to_lora_layers(
+    model: nn.Module, num_layers: int, config: Dict, use_dora: bool = ...
+):  # -> None:
+    """
+    Convert some of the models linear layers to lora layers.
+
+    Args:
+        model (nn.Module): The neural network model.
+        num_layers (int): The number of blocks to convert to lora layers
+        starting from the last layer.
+        config (dict): More configuration parameters for LoRA, including the
+          rank, scale, and optional layer keys.
+        use_dora (bool): If True, uses DoRA instead of LoRA.
+          Default: ``False``
+    """
+    ...
+
+def load_adapters(model: nn.Module, adapter_path: str) -> nn.Module:
+    """
+    Load any fine-tuned adapters / layers.
+
+    Args:
+        model (nn.Module): The neural network model.
+        adapter_path (str): Path to the adapter configuration file.
+
+    Returns:
+        nn.Module: The updated model with LoRA layers applied.
+    """
+    ...
+
+def remove_lora_layers(model: nn.Module) -> nn.Module:
+    """
+    Remove the LoRA layers from the model.
+
+    Args:
+        model (nn.Module): The model with LoRA layers.
+
+    Returns:
+        nn.Module: The model without LoRA layers.
+    """
+    ...
+
+def print_trainable_parameters(model):  # -> None:
+    ...

.mlx_typings/mlx_vlm/prompt_utils.pyi

@@ -0,0 +1,12 @@
+from typing import Any
+
+def get_message_json(
+    model_name: str,
+    prompt: str,
+    role: str = "user",
+    skip_image_token: bool = False,
+    skip_audio_token: bool = False,
+    num_images: int = 0,
+    num_audios: int = 0,
+    **kwargs: Any,
+) -> dict[str, Any]: ...

.mlx_typings/mlx_vlm/utils.pyi

@@ -0,0 +1,15 @@
+from pathlib import Path
+from typing import Any
+
+class ImageProcessor:
+    def preprocess(
+        self, images: list[dict[str, Any]], **kwargs: Any
+    ) -> dict[str, Any]: ...
+    def __call__(self, **kwargs: Any) -> dict[str, Any]: ...
+
+def load_image_processor(
+    model_path: str | Path, **kwargs: Any
+) -> ImageProcessor | None: ...
+def load_processor(
+    model_path: str | Path, add_detokenizer: bool = ..., **kwargs: Any
+) -> ImageProcessor: ...

.mlx_typings/safetensors/__init__.pyi

@@ -0,0 +1,8 @@
+from typing import Any, Self
+
+class safe_open:
+    def __init__(self, filename: str, framework: str = "pt") -> None: ...
+    def __enter__(self) -> Self: ...
+    def __exit__(self, *args: Any) -> None: ...
+    def keys(self) -> list[str]: ...
+    def get_tensor(self, name: str) -> Any: ...

.vscode/extensions.json

@@ -1,7 +1,8 @@
 {
     "recommendations": [
         "detachhead.basedpyright",
-        "ms-python.python"
+        "ms-python.python",
+        "jnoortheen.nix-ide"
     ],
     "unwantedRecommendations": [
         "ms-python.vscode-pylance",

.vscode/settings.json

@@ -1,3 +1,32 @@
 {
-    "basedpyright.importStrategy": "fromEnvironment"
+    "files.associations": {
+        "*.nix": "nix",
+    },
+    "nix.enableLanguageServer": true,
+    "nix.serverPath": "nixd",
+    "nix.serverSettings": {
+        "nixd": {
+        "formatting": {
+            "command": ["nixpkgs-fmt"]
+        },
+        "nixpkgs": {
+            "expr": "(builtins.getFlake \"path:${workspaceFolder}\").currentSystem.config._module.args.pkgs"
+        },
+        "options": {
+            "flake-parts": {
+                "expr": "(builtins.getFlake \"path:${workspaceFolder}\").debug.options"
+            },
+            "flake-parts-perSystem": {
+                "expr": "(builtins.getFlake \"path:${workspaceFolder}\").currentSystem.options"
+            }
+        }
+      }
+    },
+    "[nix]": {
+        "editor.defaultFormatter": "jnoortheen.nix-ide"
+    },
+    
+    "python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python",
+    "basedpyright.analysis.configFilePath": "${workspaceFolder}/pyproject.toml",
+    "basedpyright.importStrategy": "fromEnvironment",
 }

CONTRIBUTING.md

@@ -11,9 +11,18 @@ To run EXO from source:
   ```bash
   brew install uv
   ```
-- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+- [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
   ```bash
-  brew install macmon
+  curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
+  rustup toolchain install nightly
+  ```
+- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+  Use the pinned fork revision used by this repo instead of Homebrew `macmon`.
+  ```bash
+  cargo install --git https://github.com/vladkens/macmon \
+    --rev a1cd06b6cc0d5e61db24fd8832e74cd992097a7d \
+    macmon \
+    --force
   ```
 
 ```bash

CONVERSION_BENCH_NOTES.md

@@ -0,0 +1,482 @@
+# Tensor Conversion Benchmark Notes
+
+## Current Goal
+
+Benchmark raw tinygrad `<->` MLX tensor transformation latency on Apple Silicon
+for tensors that are already:
+
+- synchronized
+- allocated
+- materialized / realized
+
+The timed region should keep source creation and explicit synchronization
+outside the loop, while making it clear when helper, binding, owner-pinning,
+and wrapper-construction overhead are still inside it.
+
+Current benchmark CSV output reports average per-call latency plus sample
+standard deviation:
+
+- `avg_us`
+- `stddev_us`
+
+Older notes below that quote min/median values refer to earlier benchmark runs
+before the reporting format was changed.
+
+## Repo Layout
+
+- Root notes file:
+  - `CONVERSION_BENCH_NOTES.md`
+- Interop code:
+  - `mlx_tinygrad_interop/`
+- Reusable library code:
+  - `mlx_tinygrad_interop/lib/`
+- Baseline bridge module:
+  - `mlx_tinygrad_interop/lib/tensor_bridge.py`
+- Route benchmark using existing PyTorch interop:
+  - `mlx_tinygrad_interop/bench_torch_route.py`
+- Historical benchmark kept as-is:
+  - `tmp/bench_pingpong.py`
+
+## Current Fast-Path Design
+
+The first direct benchmark path is intentionally narrow.
+
+- `MLX -> tinygrad`
+  - export MLX Metal storage metadata
+  - import into tinygrad by aliasing the existing `MTLBuffer*`
+- `tinygrad -> MLX`
+  - export tinygrad Metal storage metadata
+  - import into MLX by wrapping the underlying unified-memory pointer with a
+    no-copy MLX array constructor path
+
+This is asymmetric internally:
+
+- `MLX -> tinygrad` aliases an existing `MTLBuffer*`
+- `tinygrad -> MLX` rebuilds an MLX array from a raw unified-memory pointer
+
+So the current bridge benchmark is not a symmetric measure of pure storage
+adoption cost. That is acceptable for now: the working goal is "good enough"
+bidirectional latency, not symmetry for its own sake.
+
+## Implemented Helper Surface
+
+- MLX
+  - `mx.metal._unsafe_export_storage(array)`
+  - `mx.metal._unsafe_to_tinygrad_fast(array, tg_dtype, owner=None)`
+  - `mx.metal._unsafe_rebind_tinygrad(array, borrower, owner=None)`
+  - `mx.metal._unsafe_array_from_ptr(raw_ptr, shape, dtype, owner=None)`
+  - `mx.metal._unsafe_array_from_ptr_alias_only(raw_ptr, shape, dtype, owner=None)`
+- tinygrad
+  - `Tensor._unsafe_metal_storage()`
+  - `Tensor._unsafe_from_metal_buffer(mtl_buffer_ptr, shape, dtype=..., byte_offset=0, owner=None)`
+  - `Tensor._unsafe_from_metal_buffer_fast(mtl_buffer_ptr, shape, dtype=..., byte_offset=0, owner=None)`
+  - `Tensor._unsafe_metal_borrower(mtl_buffer_ptr, shape, dtype=..., byte_offset=0, owner=None)`
+- exo handoff layer
+  - `mlx_tinygrad_interop.lib.lease_pool.MlxToTinygradLeasePool`
+  - `mlx_tinygrad_interop.lib.lease_pool.MlxToTinygradLeasePools`
+  - `mlx_tinygrad_interop.lib.lease_pool.MlxToTinygradCopyLeasePool`
+  - `mlx_tinygrad_interop.lib.lease_pool.MlxToTinygradCopyLeasePools`
+  - `mlx_tinygrad_interop.stress_interop`
+  - `mlx_tinygrad_interop.lib.tensor_bridge.tinygrad_to_mlx`
+  - `mlx_tinygrad_interop.lib.tensor_bridge.mlx_to_tinygrad` (stubbed)
+
+These helpers are intentionally private and unsafe.
+
+The current benchmark still pays Python and binding overhead in several rows.
+The older helper rows return Python dicts and unpack them before calling the
+import helper, while the newer MLX-side single-entry rows still include the
+tinygrad-side wrapper construction they trigger.
+
+## Temporary Eligibility Rules
+
+The current direct path should only accept tensors that are:
+
+- backed by Metal storage
+- already realized / available
+- single-device
+- dense row-major contiguous
+- concrete-shaped
+- dtype-compatible without conversion
+
+Non-contiguous views, broadcasts, dtype casts, and multi-device tensors should
+fall back to slower paths.
+
+For the current MLX exporter, the array must also already be in MLX's
+`available` state. The benchmark currently satisfies that with
+`mx.array(np_array)`, which is a workaround for the current helper rather than a
+claim that arbitrary lazy MLX outputs are already supported by the same path.
+
+## Workflow
+
+Use the `exo` devshell and `uv` workflow.
+
+1. Change code locally.
+2. Push the updated `mlx` and `tinygrad` fork branches.
+3. In local `exo`, enter the devshell with `nix develop`.
+4. Regenerate the lockfile against the new fork heads with:
+   `uv lock --upgrade-package mlx --refresh-package mlx --upgrade-package tinygrad --refresh-package tinygrad`
+5. Commit and push the updated `exo` branch, including the regenerated
+   `uv.lock`.
+6. On the remote Mac, pull the updated `exo` branch.
+7. Enter the devshell with `nix develop`.
+8. Refresh the environment with `uv sync`.
+9. Run tests and benchmarks with `uv run ...`.
+
+Do not rely on ad-hoc per-host build environments when the flake / devshell can
+carry the needed toolchain.
+
+For the PyTorch-route benchmark, `torch` must also be present in the normal
+Darwin project dependencies, not only in Linux extras.
+
+The current fully pre-existing PyTorch route is only usable with a CPU
+intermediate tensor on `e16`. An `mps` intermediate made the documented
+`Tensor.from_blob(..., device="METAL")` path fail when tinygrad later tried to
+use the imported object as a Metal buffer, so the route benchmark now defaults
+to `--torch-device cpu`.
+
+### Important Lockfile Note
+
+For these branch-based git dependencies, plain `uv lock` was not sufficient to
+advance the pinned SHAs in `uv.lock` during testing, and `--upgrade-package`
+alone still left a stale git revision in a later pass. The working command was:
+
+`uv lock --upgrade-package mlx --refresh-package mlx --upgrade-package tinygrad --refresh-package tinygrad`
+
+## Known Nuances / Footguns
+
+- Unified memory does not mean both frameworks consume shared storage in the
+  same way. Metal kernels still bind `MTLBuffer` objects.
+- Synchronization can dominate measured latency if it leaks into the timed path.
+- Python overhead matters at the `1-10 us` scale, so helper calls and wrapper
+  construction can dominate tiny tensors even when no tensor bytes are copied.
+- tinygrad tensors are graph objects, but once realized they do have concrete
+  underlying storage.
+- External mutation and aliasing can bypass autograd expectations in both
+  frameworks.
+- The current fast path is asymmetric:
+  - MLX exports `MTLBuffer*` for the `MLX -> tinygrad` direction.
+  - tinygrad exports raw unified-memory pointer for the `tinygrad -> MLX`
+    direction.
+- `mx.array(memoryview(...))` is not an aliasing import path in current MLX.
+  It goes through MLX's native CPU ndarray conversion path and copies the
+  bytes.
+- The first tinygrad import helper supports byte offsets.
+- The first tinygrad import helper now optionally accepts `buffer_nbytes` for
+  a bounds check. If that metadata is omitted, the helper still cannot prove
+  the requested view fits inside the borrowed buffer.
+- MLX export now distinguishes:
+  - `logical_nbytes`: the logical bytes in the exported array view
+  - `buffer_nbytes`: the backing buffer capacity
+- Offsetted MLX views must use `buffer_nbytes` semantics for bounds checks.
+- The legacy MLX export field `nbytes` was removed to avoid accidental use of
+  logical-size semantics where backing-buffer-size semantics are required.
+- tinygrad's fast Metal import helper must also treat zero-offset logical views
+  over oversized backing buffers as `BUFFER_VIEW`s. Stress testing caught a bug
+  where it tried to reshape the whole backing buffer instead.
+- MLX `buffer_nbytes` is a raw byte-capacity field, not a promise that the
+  backing buffer size is a multiple of the tensor dtype itemsize. The fast
+  tinygrad helper and borrower now use byte-level bounds checks and `ceildiv`
+  for backing-buffer sizing instead of rejecting those cases.
+- The first MLX import helper is raw-pointer based rather than foreign
+  `MTLBuffer*` based.
+- `mx.metal._unsafe_array_from_ptr(...)` may still copy if MLX cannot alias the
+  pointer directly.
+- `mx.metal._unsafe_array_from_ptr_alias_only(...)` fails instead of silently
+  copying, so it is the right helper for proving aliasing in benchmarks.
+- `mx.metal._unsafe_to_tinygrad_fast(...)` is a single MLX binding entrypoint
+  for `MLX -> tinygrad`, but it still includes tinygrad-side tensor creation.
+- `Tensor._unsafe_metal_borrower(...)` is a mutable slot primitive. It reuses
+  the same tinygrad tensor wrapper and rebinds its borrowed `MTLBuffer*`.
+  Older references to that tensor are not snapshots.
+- `borrower.rebind(...)` is now the checked path and requires explicit
+  `shape=` and `dtype_name=` compatibility. The only bypass is the private
+  `borrower._raw_rebind(...)` escape hatch kept for benchmark internals.
+- The borrower now updates its internal `external_ptr` metadata on rebind so
+  the stored buffer metadata matches the live Metal handle.
+- `mx.metal._unsafe_export_storage(...)` currently expects an MLX array that is
+  already in the C++ `available` state. In practice, `mx.array(np_array)` met
+  that precondition for local smoke testing, while `mx.arange(...)` did not.
+- The lease-pool handoff layer now owns that MLX-side availability barrier by
+  calling `mx.eval(...)` on acquire.
+- `tinygrad -> MLX memoryview_copy` also includes per-call runtime ceremony,
+  because tinygrad's zero-copy Metal memoryview export synchronizes before
+  exposing the buffer.
+- The randomized stress harness now uses `np.einsum(...)` rather than NumPy's
+  `@` operator for the `matmul_lastdim` baseline. On the current macOS
+  validation host, a valid contiguous float32 `(16,31) @ (31,7)` case through
+  NumPy `@` returned an incorrect all-zero result while MLX, tinygrad, and
+  `np.einsum` agreed on the nonzero result.
+- Bench rows named `rebindable_slot_*` measure rebindable-slot cost, not fresh
+  conversion cost.
+- Bench rows named `borrower_ring*` rotate through multiple independent slots.
+  They are intended to approximate a practical pool/ring design with fewer
+  footguns than a single slot.
+- Bench rows named `copy_pool_*` reuse tinygrad-owned destination tensors and
+  copy MLX bytes into them before release. They avoid foreign-buffer aliasing
+  but still rely on slot/pool reuse rather than fresh independent tensors.
+- The practical `MLX -> tinygrad` API shape is now a keyed lease pool:
+  - acquire from an MLX array
+  - use `lease.tensor`
+  - release the lease only after downstream work is realized / synchronized
+- That raw lease surface is still intentionally unsafe. A saved `lease.tensor`
+  reference is not a snapshot and can observe new contents if the slot is
+  later reused.
+- The preferred production-shaped handoff is now the scoped callback API:
+  - `pool.run_with_mlx_tensor(array, fn=...)`
+  - `pools.run_with_mlx_tensor(array, tg_dtype=..., fn=...)`
+  - it scopes acquire/use/release together, realizes returned tensors before
+    release, rejects returning alias views of the borrowed slot, and rejects
+    leaked live tensors whose graphs still depend on the borrowed tensor
+  - only independently realized outputs may escape the callback
+  - the callback still must not stash the raw borrowed tensor object itself;
+    that remains a contract rule rather than something the current runtime can
+    prove mechanically
+- Safe scoped release intentionally uses the global `Device["METAL"].synchronize()`
+  barrier again. The narrower callback-local command-buffer wait experiment was
+  not robust under concurrent Metal enqueue and did not show a meaningful
+  latency win in the local microbenchmarks.
+- Lease pools are keyed by `(shape, dtype, byte_offset)` so variable inference
+  shapes can be bucketed explicitly instead of silently reusing an
+  incompatible slot.
+- Both alias and copy pool registries are now bounded LRU caches with
+  `max_pools`. If the registry is full and all pools are in flight, acquire
+  fails instead of growing unbounded.
+- Bench rows named `*_then_use_sum` measure rebinding or conversion followed by
+  immediate tinygrad consumption through a realized reduction kernel.
+- Do not rebind a slot until all work derived from its previous contents has
+  been realized and synchronized. Otherwise later rebinds can change what
+  older lazy graphs or in-flight kernels observe.
+- Safe release now clears the slot's pinned `_external_owner` after the Metal
+  barrier. Unsafe `synchronize_on_release=False` flows still require the caller
+  to provide the fence discipline.
+- This optimization is same-process and same-address-space only. It does not
+  survive a process boundary or a machine boundary, and it does not remove any
+  later Metal/host -> CUDA transfer when the downstream stage runs on the RTX.
+
+## Stress Suite
+
+There is now a separate randomized stress/soak script:
+
+- `uv run python mlx_tinygrad_interop/stress_interop.py --cases 64 --soak-iterations 512`
+
+It exercises:
+
+- random shapes and dtypes
+- offsetted MLX views
+- raw conversion correctness against NumPy baselines
+- randomized downstream movement / elementwise / broadcast / reduction /
+  matmul op chains after conversion
+- roundtrip `MLX -> tinygrad -> MLX` pipeline checks after those op chains
+- repeated scoped-handoff soak loops through both alias and copy keyed pools
+- native memory reporting via:
+  - `mx.get_active_memory()`
+  - `mx.get_cache_memory()`
+  - `mx.get_peak_memory()`
+  - process `ru_maxrss`
+- bounded pool-count assertions for both alias and copy pool registries
+
+This is intended to catch value corruption, stale-slot mistakes, obvious
+crashes, and gross leak regressions before the interop path is integrated more
+deeply into the runtime.
+
+Note: the `matmul_lastdim` stress baseline uses `np.einsum(...)` instead of
+NumPy `@` because the current macOS NumPy build produced a demonstrably wrong
+all-zero result on one of the randomized float32 cases.
+
+The stress comparator is still strict, but float32 rows now use a slightly
+looser `rtol=5e-5, atol=1e-5` budget to absorb backend accumulation-order
+differences across long matmul/reduction chains without treating a few-ulps
+drift as conversion corruption.
+
+For downstream op-chain checks, the stress suite now compares against the
+native destination-framework baseline rather than NumPy directly:
+
+- `MLX -> tinygrad` post-conversion op chains are compared to a native tinygrad
+  tensor built from the same logical values
+- `tinygrad -> MLX` post-conversion op chains are compared to a native MLX
+  array built from the same logical values
+
+That split is intentional. Raw conversion is still checked against NumPy, but
+some downstream integer-promotion and reduction semantics differ between NumPy,
+MLX, and tinygrad, so the destination-framework baseline is the right
+conversion-integrity check.
+
+## Current Findings
+
+The unsafe bridge was validated through the repo-standard remote flow on `e16`:
+
+1. `git pull --ff-only`
+2. `nix develop`
+3. `uv sync`
+4. `uv run python mlx_tinygrad_interop/bench_raw_conversion.py ...`
+
+The direct helpers worked in both directions:
+
+- `MLX -> tinygrad` unsafe helper bridge returned correct values.
+- `tinygrad -> MLX` unsafe helper bridge returned correct values.
+- The expanded slot / ring tests passed on `e16`:
+  - nonzero-offset MLX slice import
+  - single-entry `MLX -> tinygrad`
+  - mutable-slot rebinding semantics
+  - shape/dtype mismatch rejection on rebind
+  - slot metadata update on rebind
+  - two-slot independence until a slot is reused
+- After adding shape/dtype contract checks, a spot-check rerun at `7168` bytes
+  showed:
+  - `rebindable_slot_bridge`: `2.458 us` min, `2.494 us` median
+  - `borrower_ring4_bridge`: `2.545 us` min, `2.566 us` median
+  - `rebindable_slot_import_only`: `2.395 us` min, `2.400 us` median
+  - `borrower_ring4_import_only`: `2.449 us` min, `2.477 us` median
+- So the contract hardening added roughly `~1 us`, but the slot/ring path still
+  remains comfortably inside the target latency class.
+
+Updated remote latency measurements for `float32` and `7168` bytes were:
+
+- `unsafe_helper_bridge`
+  - `mlx_to_tinygrad`: `21.202 us` min, `21.532 us` median
+  - `tinygrad_to_mlx`: `28.109 us` min, `28.372 us` median
+- `single_entry_bridge`
+  - `mlx_to_tinygrad`: `21.388 us` min, `21.542 us` median
+- `fresh_wrapper_then_use_sum`
+  - `mlx_to_tinygrad`: `601.812 us` min, `611.458 us` median
+- `rebindable_slot_bridge`
+  - `mlx_to_tinygrad`: `1.505 us` min, `1.542 us` median
+- `rebindable_slot_then_use_sum`
+  - `mlx_to_tinygrad`: `579.583 us` min, `581.833 us` median
+- `borrower_ring4_bridge`
+  - `mlx_to_tinygrad`: `1.531 us` min, `1.573 us` median
+- `borrower_ring4_then_use_sum`
+  - `mlx_to_tinygrad`: `577.730 us` min, `581.000 us` median
+- `unsafe_helper_legacy`
+  - `mlx_to_tinygrad`: `31.938 us` min, `32.214 us` median
+- `unsafe_helper_maybe_copy`
+  - `tinygrad_to_mlx`: `28.153 us` min, `28.277 us` median
+- `memoryview_copy`
+  - `mlx_to_tinygrad`: `35.191 us` min, `35.668 us` median
+  - `tinygrad_to_mlx`: `2.596 us` min, `2.662 us` median
+- `numpy_baseline`
+  - `mlx_to_tinygrad`: `272.323 us` min, `275.104 us` median
+  - `tinygrad_to_mlx`: `12.817 us` min, `13.005 us` median
+
+Interpretation:
+
+- The lower-overhead tinygrad import helper cut `MLX -> tinygrad` from about
+  `32 us` to about `21 us` at `7 kB`, so the old `Tensor.empty(...)` based
+  helper was a real source of overhead.
+- Replacing the exporter dict/unpack stack with a single MLX binding entrypoint
+  still barely moved `MLX -> tinygrad`: about `21.2 us -> 21.4 us` at `7 kB`.
+- That means the remaining fixed cost was not materially in MLX export or
+  Python exporter marshalling. It was overwhelmingly on the tinygrad side.
+- Rebinding a single mutable slot dropped `MLX -> tinygrad` to about `1.5 us`
+  at `7 kB`. Rotating through a ring of four slots landed at about `1.53 us`,
+  so the ring preserved essentially all of the latency win.
+- Those slot / ring rows are not fresh conversion results. They are
+  rebind-and-return-slot results.
+- That is still inside the target range and strongly indicates that tinygrad
+  wrapper construction, not storage adoption itself, was the dominant cost.
+- On `e16`, the strict alias-only `tinygrad -> MLX` helper succeeded. Its
+  timings were effectively the same as the maybe-copy helper, so the benchmark
+  can now report a proven aliasing path in that direction on this host.
+- `tinygrad -> MLX` currently has a very cheap copy path because `mx.array()`
+  over a tinygrad `memoryview` is implemented efficiently in MLX's native C++
+  import path, even though it still copies.
+- At this tensor size, Python call overhead and wrapper construction matter
+  much more than raw byte movement.
+- The `*_then_use_sum` rows are dominated by the tinygrad reduction kernel
+  itself. They are useful as end-to-end "convert then immediately consume"
+  probes, not as pure conversion timings.
+- Those end-to-end rows still show the same direction: at `7 kB`, slot / ring
+  rebinding saves roughly `20-25 us` versus the fresh-wrapper path, but the
+  absolute runtime is around `580-600 us` because the reduction dominates.
+- These numbers do not establish that "aliasing costs ~21-28 us". They
+  establish that creating a fresh tinygrad wrapper through the current helper
+  stack costs that much, while rebinding a pre-existing slot costs about
+  `1.5 us` on this host.
+- An offsetted MLX slice was also validated through the new export semantics:
+  `offset_bytes=64`, `logical_nbytes=7168`, `buffer_nbytes=16384`, and the
+  borrowed tinygrad tensor matched the expected values.
+- The rebindable slot is narrower than a normal conversion helper:
+  - it returns the same tinygrad `Tensor` object each time
+  - it assumes fixed shape / dtype / byte-offset semantics
+  - older references are not snapshots
+  - safe reuse requires that all work from the previous lease has already been
+    realized and synchronized
+  - it is therefore best understood as a dangerous but very informative lower
+    bound and a candidate building block for a specialized converter API
+- A ring of multiple slots is the more practical extension of that idea because
+  it preserves most of the latency win while reducing the worst single-slot
+  footgun.
+
+Additional remote microbench sweep on `e16` for `256`, `7168`, `65536`, and
+`1048576` bytes showed:
+
+- `MLX -> tinygrad`
+  - `unsafe_helper_bridge`: roughly `21-23 us`
+  - `single_entry_bridge`: roughly `21-22 us`
+  - `fresh_wrapper_then_use_sum`: roughly `601-693 us`
+  - `rebindable_slot_bridge`: roughly `1.49-1.53 us`
+  - `rebindable_slot_then_use_sum`: roughly `573-670 us`
+  - `borrower_ring4_bridge`: roughly `1.53-1.62 us`
+  - `borrower_ring4_then_use_sum`: roughly `566-669 us`
+  - `unsafe_helper_legacy`: roughly `31-33 us`
+  - `memoryview_copy`: roughly `34-51 us`
+  - `numpy_baseline`: roughly `269-309 us`
+  - `export_helper_only`: roughly `0.58-0.63 us`
+  - `import_helper_fast_only`: roughly `20.2-21.4 us`
+  - `rebindable_slot_import_only`: roughly `1.40-1.47 us`
+  - `borrower_ring4_import_only`: roughly `1.45-1.49 us`
+  - `import_helper_legacy_only`: roughly `31-33 us`
+- `tinygrad -> MLX`
+  - `unsafe_helper_bridge`: roughly `27-29 us`
+  - `unsafe_helper_maybe_copy`: roughly `27-29 us`
+  - `memoryview_copy`: roughly `2.5 us` at `256 B`, `2.5 us` at `7168 B`,
+    `3.5 us` at `64 KiB`, and `17.4 us` at `1 MiB`
+  - `numpy_baseline`: roughly `12.7 us` at `256 B`, `12.7 us` at `7168 B`,
+    `16.6 us` at `64 KiB`, and `56.7 us` at `1 MiB`
+  - `export_helper_only`: roughly `23.1-24.4 us`
+  - `import_helper_only`: roughly `2.12-2.22 us`
+  - `import_helper_maybe_copy_only`: roughly `2.16-2.29 us`
+
+What this means:
+
+- The MLX exporter is already cheap, and even a single MLX binding entrypoint
+  did not change `MLX -> tinygrad` materially. That closes out the
+  "Python exporter ceremony" hypothesis for the current bridge.
+- The MLX importer from raw pointer is also already cheap, whether measured in
+  strict alias-only mode or maybe-copy mode on this host.
+- The lower-overhead tinygrad import helper bought a real speedup, but the
+  expensive piece for `MLX -> tinygrad` was still constructing a fresh tinygrad
+  wrapper around the borrowed storage.
+- Rebinding a slot or rotating through a ring changes the latency class
+  completely. The "rebind one pre-existing slot" lower bound is about
+  `1.5 us` on this host for the measured sizes, and a ring of four slots keeps
+  essentially the same latency.
+- For `tinygrad -> MLX`, the native copy path is already in the desired latency
+  class for small tensors and remains competitive well past `7 kB`.
+- For `MLX -> tinygrad`, a fresh-wrapper helper is still not close to the
+  desired `1-10 us` range at `7 kB`, but a slot / ring helper is.
+- For end-to-end "convert then immediately use" measurements, the tinygrad
+  compute dominates. The relevant signal is the delta versus the fresh-wrapper
+  path, not the absolute `~580-700 us` number.
+
+## Near-Term Plan
+
+1. Treat `tinygrad -> MLX memoryview_copy` as the current practical fast path.
+2. Treat `MLX -> tinygrad` slot/ring rebinding as the current latency floor and
+   likely practical fast path when rebindable slot semantics are acceptable.
+3. Prefer a ring/pool of slots over a single slot for any practical design.
+4. If `MLX -> tinygrad` must return a fresh tinygrad tensor each time and still
+   stay under `10 us`, the remaining work is entirely on the tinygrad-side
+   construction path.
+5. Avoid spending time on symmetry unless it becomes necessary for a specific
+   downstream use case.
+6. If this moves into the real MLX/tinygrad disaggregated runtime, the next API
+   shape should be an explicit pool/lease abstraction with generation tracking,
+   not bare mutable-slot rebinding.
+
+## Open Questions
+
+- Whether the first fast path should support contiguous slices with byte
+  offsets, or only base-contiguous tensors.
+- Whether the slot/ring primitive should stay benchmark-only or be surfaced as
+  a deliberate specialized converter API.

README.md

@@ -95,11 +95,10 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
   /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
   ```
 - [uv](https://github.com/astral-sh/uv) (for Python dependency management)
-- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
 - [node](https://github.com/nodejs/node) (for building the dashboard)
 
   ```bash
-  brew install uv macmon node
+  brew install uv node
   ```
 - [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
 
@@ -107,6 +106,17 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
   curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
   rustup toolchain install nightly
   ```
+- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+
+  Install the pinned fork revision used by this repo instead of Homebrew `macmon`.
+  Homebrew `macmon 0.6.1` still crashes on Apple M5.
+
+  ```bash
+  cargo install --git https://github.com/vladkens/macmon \
+    --rev a1cd06b6cc0d5e61db24fd8832e74cd992097a7d \
+    macmon \
+    --force
+  ```
 
 Clone the repo, build the dashboard, and run exo:
 
@@ -285,8 +295,9 @@ exo supports several environment variables for configuration:
 
 | Variable | Description | Default |
 |----------|-------------|---------|
-| `EXO_MODELS_PATH` | Colon-separated paths to search for pre-downloaded models (e.g., on NFS mounts or shared storage) | None |
-| `EXO_MODELS_DIR` | Directory where exo downloads and stores models | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
+| `EXO_DEFAULT_MODELS_DIR` | Default directory for model downloads and caches. Always first in the writable dirs list. | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
+| `EXO_MODELS_DIRS` | Colon-separated additional writable directories for model downloads. Checked in order after the default; first with enough free space is used. | None |
+| `EXO_MODELS_READ_ONLY_DIRS` | Colon-separated read-only directories to search for pre-downloaded models (e.g., NFS mounts, shared storage). Models here cannot be deleted. | None |
 | `EXO_OFFLINE` | Run without internet connection (uses only local models) | `false` |
 | `EXO_ENABLE_IMAGE_MODELS` | Enable image model support | `false` |
 | `EXO_LIBP2P_NAMESPACE` | Custom namespace for cluster isolation | None |
@@ -296,8 +307,11 @@ exo supports several environment variables for configuration:
 **Example usage:**
 
 ```bash
-# Use pre-downloaded models from NFS mount
-EXO_MODELS_PATH=/mnt/nfs/models:/opt/ai-models uv run exo
+# Use pre-downloaded models from NFS mount (read-only)
+EXO_MODELS_READ_ONLY_DIRS=/mnt/nfs/models:/opt/ai-models uv run exo
+
+# Download models to an external SSD (falls back to default dir if full)
+EXO_MODELS_DIRS=/Volumes/ExternalSSD/exo-models uv run exo
 
 # Run in offline mode
 EXO_OFFLINE=true uv run exo

app/EXO/EXO/ExoProcessController.swift

@@ -4,9 +4,30 @@ import Foundation
 
 private let customNamespaceKey = "EXOCustomNamespace"
 private let hfTokenKey = "EXOHFToken"
+private let hfEndpointKey = "EXOHFEndpoint"
 private let enableImageModelsKey = "EXOEnableImageModels"
 private let offlineModeKey = "EXOOfflineMode"
+private let fastSynchEnabledKey = "EXOFastSynchEnabled"
 private let onboardingCompletedKey = "EXOOnboardingCompleted"
+private let defaultModelsDirKey = "EXODefaultModelsDir"
+private let additionalModelsDirsKey = "EXOAdditionalModelsDirs"
+private let readOnlyModelsDirsKey = "EXOReadOnlyModelsDirs"
+private let customEnvironmentVariablesKey = "EXOCustomEnvironmentVariables"
+
+/// A user-defined environment variable that is injected into the exo child
+/// process at launch. Used as an escape hatch for env vars that don't have
+/// first-class typed UI in Settings.
+struct CustomEnvironmentVariable: Codable, Identifiable, Equatable {
+    var id: UUID
+    var key: String
+    var value: String
+
+    init(id: UUID = UUID(), key: String = "", value: String = "") {
+        self.id = id
+        self.key = key
+        self.value = value
+    }
+}
 
 @MainActor
 final class ExoProcessController: ObservableObject {
@@ -53,6 +74,14 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(hfToken, forKey: hfTokenKey)
         }
     }
+    @Published var hfEndpoint: String = {
+        return UserDefaults.standard.string(forKey: hfEndpointKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(hfEndpoint, forKey: hfEndpointKey)
+        }
+    }
     @Published var enableImageModels: Bool = {
         return UserDefaults.standard.bool(forKey: enableImageModelsKey)
     }()
@@ -69,6 +98,60 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
         }
     }
+    @Published var fastSynchEnabled: Bool = {
+        if UserDefaults.standard.object(forKey: fastSynchEnabledKey) == nil {
+            return true
+        }
+        return UserDefaults.standard.bool(forKey: fastSynchEnabledKey)
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(fastSynchEnabled, forKey: fastSynchEnabledKey)
+        }
+    }
+    @Published var defaultModelsDir: String = {
+        return UserDefaults.standard.string(forKey: defaultModelsDirKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(defaultModelsDir, forKey: defaultModelsDirKey)
+        }
+    }
+    @Published var additionalModelsDirs: String = {
+        return UserDefaults.standard.string(forKey: additionalModelsDirsKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(additionalModelsDirs, forKey: additionalModelsDirsKey)
+        }
+    }
+    @Published var readOnlyModelsDirs: String = {
+        return UserDefaults.standard.string(forKey: readOnlyModelsDirsKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(readOnlyModelsDirs, forKey: readOnlyModelsDirsKey)
+        }
+    }
+    @Published var customEnvironmentVariables: [CustomEnvironmentVariable] = {
+        guard
+            let data = UserDefaults.standard.data(forKey: customEnvironmentVariablesKey),
+            let decoded = try? JSONDecoder().decode(
+                [CustomEnvironmentVariable].self, from: data
+            )
+        else {
+            return []
+        }
+        return decoded
+    }()
+    {
+        didSet {
+            guard let data = try? JSONEncoder().encode(customEnvironmentVariables) else {
+                return
+            }
+            UserDefaults.standard.set(data, forKey: customEnvironmentVariablesKey)
+        }
+    }
 
     /// Fires once when EXO transitions to `.running` for the very first time (fresh install).
     @Published private(set) var isFirstLaunchReady = false
@@ -273,12 +356,16 @@ final class ExoProcessController: ObservableObject {
         if !hfToken.isEmpty {
             environment["HF_TOKEN"] = hfToken
         }
+        if !hfEndpoint.isEmpty {
+            environment["HF_ENDPOINT"] = hfEndpoint
+        }
         if enableImageModels {
             environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
         }
         if offlineMode {
             environment["EXO_OFFLINE"] = "true"
         }
+        environment["EXO_FAST_SYNCH"] = fastSynchEnabled ? "true" : "false"
 
         var paths: [String] = []
         if let existing = environment["PATH"], !existing.isEmpty {
@@ -303,6 +390,29 @@ final class ExoProcessController: ObservableObject {
         }
 
         environment["PATH"] = paths.joined(separator: ":")
+
+        let trimmedDefaultModelsDir = defaultModelsDir.trimmingCharacters(in: .whitespaces)
+        if !trimmedDefaultModelsDir.isEmpty {
+            environment["EXO_DEFAULT_MODELS_DIR"] = trimmedDefaultModelsDir
+        }
+        let trimmedAdditionalModelsDirs = additionalModelsDirs.trimmingCharacters(in: .whitespaces)
+        if !trimmedAdditionalModelsDirs.isEmpty {
+            environment["EXO_MODELS_DIRS"] = trimmedAdditionalModelsDirs
+        }
+        let trimmedReadOnlyModelsDirs = readOnlyModelsDirs.trimmingCharacters(in: .whitespaces)
+        if !trimmedReadOnlyModelsDirs.isEmpty {
+            environment["EXO_MODELS_READ_ONLY_DIRS"] = trimmedReadOnlyModelsDirs
+        }
+
+        // Apply user-defined arbitrary environment variables last so that
+        // power users can override any of the typed fields above when
+        // necessary. Empty keys are ignored.
+        for variable in customEnvironmentVariables {
+            let trimmedKey = variable.key.trimmingCharacters(in: .whitespaces)
+            guard !trimmedKey.isEmpty else { continue }
+            environment[trimmedKey] = variable.value
+        }
+
         return environment
     }
 

app/EXO/EXO/Views/SettingsView.swift

@@ -12,8 +12,14 @@ struct SettingsView: View {
 
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
+    @State private var pendingHFEndpoint: String = ""
     @State private var pendingEnableImageModels = false
     @State private var pendingOfflineMode = false
+    @State private var pendingFastSynchEnabled = false
+    @State private var pendingDefaultModelsDir: String = ""
+    @State private var pendingAdditionalModelsDirs: String = ""
+    @State private var pendingReadOnlyModelsDirs: String = ""
+    @State private var pendingCustomEnvironmentVariables: [CustomEnvironmentVariable] = []
     @State private var needsRestart = false
     @State private var bugReportInFlight = false
     @State private var bugReportMessage: String?
@@ -33,17 +39,27 @@ struct SettingsView: View {
                 .tabItem {
                     Label("Advanced", systemImage: "wrench.and.screwdriver")
                 }
+            environmentTab
+                .tabItem {
+                    Label("Environment", systemImage: "terminal")
+                }
             aboutTab
                 .tabItem {
                     Label("About", systemImage: "info.circle")
                 }
         }
-        .frame(width: 450, height: 400)
+        .frame(width: 640, height: 560)
         .onAppear {
             pendingNamespace = controller.customNamespace
             pendingHFToken = controller.hfToken
+            pendingHFEndpoint = controller.hfEndpoint
             pendingEnableImageModels = controller.enableImageModels
             pendingOfflineMode = controller.offlineMode
+            pendingFastSynchEnabled = controller.fastSynchEnabled
+            pendingDefaultModelsDir = controller.defaultModelsDir
+            pendingAdditionalModelsDirs = controller.additionalModelsDirs
+            pendingReadOnlyModelsDirs = controller.readOnlyModelsDirs
+            pendingCustomEnvironmentVariables = controller.customEnvironmentVariables
             needsRestart = false
         }
     }
@@ -54,9 +70,9 @@ struct SettingsView: View {
         Form {
             Section {
                 LabeledContent("Cluster Namespace") {
-                    TextField("default", text: $pendingNamespace)
+                    TextField("", text: $pendingNamespace, prompt: Text("default"))
                         .textFieldStyle(.roundedBorder)
-                        .frame(width: 200)
+                        .frame(width: 260)
                 }
                 Text("Nodes with the same namespace form a cluster. Leave empty for default.")
                     .font(.caption)
@@ -65,15 +81,26 @@ struct SettingsView: View {
 
             Section {
                 LabeledContent("HuggingFace Token") {
-                    SecureField("optional", text: $pendingHFToken)
+                    SecureField("", text: $pendingHFToken, prompt: Text("optional"))
                         .textFieldStyle(.roundedBorder)
-                        .frame(width: 200)
+                        .frame(width: 260)
                 }
                 Text("Required for gated models. Get yours at huggingface.co/settings/tokens")
                     .font(.caption)
                     .foregroundColor(.secondary)
             }
 
+            Section {
+                LabeledContent("HuggingFace Endpoint") {
+                    TextField("", text: $pendingHFEndpoint, prompt: Text("default"))
+                        .textFieldStyle(.roundedBorder)
+                        .frame(width: 260)
+                }
+                Text("Defaults to huggingface.co. Use a mirror (e.g. hf-mirror.com) for China.")
+                    .font(.caption)
+                    .foregroundColor(.secondary)
+            }
+
             Section {
                 Toggle("Offline Mode", isOn: $pendingOfflineMode)
                 Text("Skip internet checks and use only locally available models.")
@@ -124,6 +151,23 @@ struct SettingsView: View {
 
     private var advancedTab: some View {
         Form {
+            Section("Performance") {
+                Toggle("Fast Synch Enabled", isOn: $pendingFastSynchEnabled)
+                Text(
+                    "Experimental: enables fast CPU to GPU synchronization. Can sometimes cause a \"GPU lock\" where inference hangs for ~10 seconds before starting. Necessary for low latency with RDMA and Tensor Parallelism."
+                )
+                .font(.caption)
+                .foregroundColor(.secondary)
+
+                HStack {
+                    Spacer()
+                    Button("Save & Restart") {
+                        applyAdvancedSettings()
+                    }
+                    .disabled(!hasAdvancedChanges)
+                }
+            }
+
             Section("Onboarding") {
                 HStack {
                     VStack(alignment: .leading) {
@@ -180,6 +224,128 @@ struct SettingsView: View {
         .padding()
     }
 
+    // MARK: - Environment Tab
+
+    private var environmentTab: some View {
+        Form {
+            Section("Models Directories") {
+                LabeledContent("Default Models Directory") {
+                    TextField(
+                        "",
+                        text: $pendingDefaultModelsDir,
+                        prompt: Text("~/.exo/models")
+                    )
+                    .textFieldStyle(.roundedBorder)
+                    .font(.system(.body, design: .monospaced))
+                    .frame(width: 260)
+                }
+                Text("Sets EXO_DEFAULT_MODELS_DIR. Where models are downloaded.")
+                    .font(.caption)
+                    .foregroundColor(.secondary)
+
+                LabeledContent("Additional Directories") {
+                    TextField(
+                        "",
+                        text: $pendingAdditionalModelsDirs,
+                        prompt: Text("optional, colon-separated")
+                    )
+                    .textFieldStyle(.roundedBorder)
+                    .font(.system(.body, design: .monospaced))
+                    .frame(width: 260)
+                }
+                Text("Sets EXO_MODELS_DIRS. Extra writable model directories.")
+                    .font(.caption)
+                    .foregroundColor(.secondary)
+
+                LabeledContent("Read-Only Directories") {
+                    TextField(
+                        "",
+                        text: $pendingReadOnlyModelsDirs,
+                        prompt: Text("optional, colon-separated")
+                    )
+                    .textFieldStyle(.roundedBorder)
+                    .font(.system(.body, design: .monospaced))
+                    .frame(width: 260)
+                }
+                Text("Sets EXO_MODELS_READ_ONLY_DIRS. Never written to.")
+                    .font(.caption)
+                    .foregroundColor(.secondary)
+            }
+
+            Section("Custom Environment Variables") {
+                Text(
+                    "Escape hatch for env vars that don't have typed fields above. "
+                        + "Values here override the typed fields on conflict."
+                )
+                .font(.caption)
+                .foregroundColor(.secondary)
+
+                if pendingCustomEnvironmentVariables.isEmpty {
+                    Text("No custom variables.")
+                        .font(.caption)
+                        .foregroundColor(.secondary)
+                } else {
+                    ForEach($pendingCustomEnvironmentVariables) { $variable in
+                        HStack(alignment: .center, spacing: 8) {
+                            VStack(spacing: 4) {
+                                TextField("key", text: $variable.key)
+                                    .labelsHidden()
+                                    .textFieldStyle(.roundedBorder)
+                                    .font(.system(.body, design: .monospaced))
+                                TextField("value", text: $variable.value)
+                                    .labelsHidden()
+                                    .textFieldStyle(.roundedBorder)
+                                    .font(.system(.body, design: .monospaced))
+                            }
+                            VStack(spacing: 4) {
+                                Button {
+                                    pendingCustomEnvironmentVariables.removeAll {
+                                        $0.id == variable.id
+                                    }
+                                } label: {
+                                    Image(systemName: "minus.circle")
+                                }
+                                .buttonStyle(.borderless)
+                                .help("Remove variable")
+                                if !isValidEnvironmentVariableName(variable.key) {
+                                    Image(systemName: "exclamationmark.triangle.fill")
+                                        .foregroundColor(.orange)
+                                        .help(
+                                            "Invalid environment variable name. "
+                                                + "Must match [A-Za-z_][A-Za-z0-9_]*."
+                                        )
+                                }
+                            }
+                        }
+                    }
+                }
+
+                HStack {
+                    Button {
+                        pendingCustomEnvironmentVariables.append(
+                            CustomEnvironmentVariable()
+                        )
+                    } label: {
+                        Label("Add Variable", systemImage: "plus")
+                    }
+                    Spacer()
+                }
+            }
+
+            Section {
+                HStack {
+                    Spacer()
+                    Button("Save & Restart") {
+                        applyEnvironmentSettings()
+                    }
+                    .disabled(!hasEnvironmentChanges)
+                }
+            }
+        }
+        .formStyle(.grouped)
+        .padding()
+    }
+
     // MARK: - About Tab
 
     private var aboutTab: some View {
@@ -454,6 +620,7 @@ struct SettingsView: View {
 
     private var hasGeneralChanges: Bool {
         pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
+            || pendingHFEndpoint != controller.hfEndpoint
             || pendingOfflineMode != controller.offlineMode
     }
 
@@ -461,9 +628,21 @@ struct SettingsView: View {
         pendingEnableImageModels != controller.enableImageModels
     }
 
+    private var hasAdvancedChanges: Bool {
+        pendingFastSynchEnabled != controller.fastSynchEnabled
+    }
+
+    private var hasEnvironmentChanges: Bool {
+        pendingDefaultModelsDir != controller.defaultModelsDir
+            || pendingAdditionalModelsDirs != controller.additionalModelsDirs
+            || pendingReadOnlyModelsDirs != controller.readOnlyModelsDirs
+            || pendingCustomEnvironmentVariables != controller.customEnvironmentVariables
+    }
+
     private func applyGeneralSettings() {
         controller.customNamespace = pendingNamespace
         controller.hfToken = pendingHFToken
+        controller.hfEndpoint = pendingHFEndpoint
         controller.offlineMode = pendingOfflineMode
         restartIfRunning()
     }
@@ -473,6 +652,75 @@ struct SettingsView: View {
         restartIfRunning()
     }
 
+    private func applyAdvancedSettings() {
+        controller.fastSynchEnabled = pendingFastSynchEnabled
+        restartIfRunning()
+    }
+
+    private func applyEnvironmentSettings() {
+        controller.defaultModelsDir = pendingDefaultModelsDir.trimmingCharacters(
+            in: .whitespaces)
+        controller.additionalModelsDirs = pendingAdditionalModelsDirs.trimmingCharacters(
+            in: .whitespaces)
+        controller.readOnlyModelsDirs = pendingReadOnlyModelsDirs.trimmingCharacters(
+            in: .whitespaces)
+
+        pendingDefaultModelsDir = controller.defaultModelsDir
+        pendingAdditionalModelsDirs = controller.additionalModelsDirs
+        pendingReadOnlyModelsDirs = controller.readOnlyModelsDirs
+
+        // Trim whitespace from keys and drop empty ones so that the stored
+        // form matches what is actually injected into the child process and
+        // hasEnvironmentChanges doesn't show a stale diff after save.
+        let trimmed: [CustomEnvironmentVariable] =
+            pendingCustomEnvironmentVariables.compactMap { variable in
+                let key = variable.key.trimmingCharacters(in: .whitespaces)
+                guard !key.isEmpty else { return nil }
+                return CustomEnvironmentVariable(
+                    id: variable.id, key: key, value: variable.value
+                )
+            }
+
+        // De-duplicate keys, keeping the last occurrence. This matches the
+        // effective semantics of the dictionary assignment in
+        // ExoProcessController.makeEnvironment and avoids silently losing
+        // visible rows after save.
+        var seenKeys = Set<String>()
+        var deduplicatedReversed: [CustomEnvironmentVariable] = []
+        for variable in trimmed.reversed() {
+            if seenKeys.insert(variable.key).inserted {
+                deduplicatedReversed.append(variable)
+            }
+        }
+        let sanitized = Array(deduplicatedReversed.reversed())
+
+        pendingCustomEnvironmentVariables = sanitized
+        controller.customEnvironmentVariables = sanitized
+
+        restartIfRunning()
+    }
+
+    /// Validates a POSIX-style environment variable name:
+    /// `[A-Za-z_][A-Za-z0-9_]*`. Uses an ASCII-only charset so that
+    /// Unicode letters (e.g. `ñ`, Cyrillic) are rejected in line with what
+    /// the help tooltip advertises. Empty strings are treated as valid
+    /// here so that a freshly added blank row does not immediately look
+    /// broken; the save step filters empty keys out instead.
+    private func isValidEnvironmentVariableName(_ key: String) -> Bool {
+        if key.isEmpty { return true }
+        let headAllowed = CharacterSet(
+            charactersIn: "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz_"
+        )
+        let tailAllowed = headAllowed.union(CharacterSet(charactersIn: "0123456789"))
+        guard let first = key.unicodeScalars.first, headAllowed.contains(first) else {
+            return false
+        }
+        for scalar in key.unicodeScalars.dropFirst() {
+            if !tailAllowed.contains(scalar) { return false }
+        }
+        return true
+    }
+
     private func restartIfRunning() {
         if controller.status == .running || controller.status == .starting {
             controller.restart()

app/EXO/EXO/Views/SettingsWindowController.swift

@@ -30,7 +30,7 @@ final class SettingsWindowController: ObservableObject {
         let hostingView = NSHostingView(rootView: settingsView)
 
         let newWindow = NSWindow(
-            contentRect: NSRect(x: 0, y: 0, width: 450, height: 400),
+            contentRect: NSRect(x: 0, y: 0, width: 640, height: 560),
             styleMask: [.titled, .closable],
             backing: .buffered,
             defer: false

bench/METHODOLOGY.md

@@ -0,0 +1,160 @@
+# Exo-Bench — Methodology
+
+exo bench measures inference throughput and resource consumption of an exo cluster under controlled conditions. It sends prompts to the `/bench/chat/completions` endpoint, collects server-reported timing statistics, and records system-level metrics (power, GPU utilisation, temperature) throughout each run.
+
+The goal is to have accurate, transparent and reproducible numbers to compare speed and scaling across different models and different setups, and to be able to track these results as optimizations and features are added to EXO.
+
+Below is a technical summary of how Exo-Bench works. While the methodology and benchmark may change over time, this document will be kept up to date whenever this happens. If you find an issue with the methodology, or would like a  feature to be added, please open a GitHub issue!
+
+---
+
+## Prompt Construction
+
+Benchmarks need prompts of an exact token length. Unfortunately, we do not have direct access to the model but just the chat completion endpoint. To get around this fact, we create a request that will tokenise to a certain prompt length.
+
+This is achieved by:
+
+1. Tokenising a sample message through the model's `apply_chat_template()` to measure overhead (system tokens, special tokens, chat formatting).
+2. Binary-searching over a repeated atom string (default `"a "`) to find the content length that produces exactly the target number of tokens after template expansion.
+3. Returning both the content string and the verified token count.
+
+The actual token count is recorded in every result row as `pp_tokens`, so downstream analysis can confirm the prompt hit its target.
+
+Chat template formatting means that it may be impossible to attain very small pp benchmarks. e.g. pp=32 may not work. This tradeoff was made because the result of such a small prompt does not seem very interesting or useful for any real-world use cases.
+
+---
+
+## Bench Endpoint
+
+When a request reaches the server via the `/bench/chat/completions` endpoint, three things change compared to a normal chat completion:
+
+- **KV prefix cache is disabled by default**. Every request starts from a cold cache, ensuring prefill timing is not affected by prior requests. See [Prefix Cache Mode](#prefix-cache-mode) for the `--use-prefix-cache` option.
+- **EOS tokens are banned**. A logits processor suppresses all end-of-sequence tokens, forcing the model to generate exactly `max_tokens` tokens. This guarantees consistent generation length for fair TPS comparison — the model cannot short-circuit a run by stopping early.
+- **No model output parsing**. The bench collection path concatenates raw token text without any model-specific post-processing (thinking tag extraction, structured output handling, etc.). This is to avoid model outputs such as tool parsing or any structural mistakes from breaking the benchmark - we are testing for speed; see Exo-Eval for performance metrics.
+
+---
+
+## Timing
+
+### Prefill TPS
+
+Measured server-side per task.
+
+```
+prefill_tps = num_prompt_tokens / prefill_wall_seconds
+```
+
+### Generation TPS
+
+Measured server-side per task. Each task records wall-clock timestamps as tokens arrive:
+
+- First generated token: timestamp recorded
+- Every subsequent token: timestamp updated
+
+When generation completes:
+
+```
+gen_span = last_token_time - first_token_time
+generation_tps = (completion_tokens - 1) / gen_span
+```
+
+The first token is excluded from the numerator because the rate measures inter-token throughput — the time between the first and last token divided by the number of intervals.
+
+This does mean that tg=1 will not work.
+
+---
+
+## Concurrency
+
+### Single Request
+
+The client records wall-clock `elapsed_s` around the HTTP round-trip (network latency + server prefill + generation + response serialisation). This is a convenience metric for end-to-end latency. The authoritative TPS numbers come from the server-side per-task timing in the `generation_stats` response.
+
+### Concurrent Requests
+
+When `--concurrency N` is set with N > 1, all N requests must hit the server at the same instant. The mechanism:
+
+1. The prompt is built once and shared across all threads.
+2. Each thread gets its own HTTP connection.
+3. A thread barrier blocks all threads until every thread is ready.
+4. The first thread past the barrier records the batch start time and signals the others.
+5. All threads use the same start time as their reference, then fire their HTTP request.
+6. Each thread's `elapsed_s` is measured from the shared start time to its own response completion.
+
+**Batch wall time** is the maximum `elapsed_s` across all N requests — the time until the last request finishes.
+
+### Aggregate TPS
+
+```
+per_req_tps = max(generation_tps across N concurrent requests)
+agg_gen_tps = per_req_tps * concurrency
+```
+
+`max` is used instead of `mean` because all requests run in parallel against the same model. The fastest request's generation rate represents the system's per-stream throughput capacity; multiplying by concurrency gives aggregate throughput.
+
+---
+
+## Prefix Cache Mode
+
+When `--use-prefix-cache` is passed, the KV prefix cache remains active during benchmarking. This speeds up repeated runs by skipping redundant prefill work, which is useful when prompt processing is not the focus of the benchmark (e.g. when measuring generation throughput or power consumption across many configurations).
+
+Each response includes a `prefix_cache_hit` field (`"none"`, `"partial"`, or `"exact"`):
+
+- **none**: Cold prefill — no cached KV state was available. The reported `prompt_tps` is the real prefill throughput.
+- **partial**: A prefix of the prompt was found in cache. Only the remaining tokens were prefilled. The reported `prompt_tps` reflects the real throughput on the uncached portion. This occurs when multiple ascending `--pp` values share a common prefix (e.g. `--pp 1000,5000` — the 5000-token prompt reuses the 1000-token cache entry and prefills the remaining 4000 tokens).
+- **exact**: The entire prompt was found in cache (e.g. same `--pp` value on a `--repeat`). No prefill work was done. The reported `prompt_tps` is the TPS from when the cache entry was originally created, not a new measurement.
+
+**Prompt TPS is approximate in this mode.** Exact-hit runs report the stored TPS from the original cold/partial prefill rather than a freshly measured value. For accurate cold prefill numbers, run without `--use-prefix-cache`.
+
+Ascending `--pp` order (e.g. `--pp 1000,5000,10000`) gives the most useful data: each size gets a meaningful partial hit except the first which is cold. Descending order produces exact hits with approximate TPS from a longer prompt's original run.
+
+---
+
+## Warmup
+
+Before measurement begins, `--warmup N` (default: 0) discarded requests are sent using the first pp/tg pair. Warmup results are not included in the output.
+
+---
+
+## System Metrics
+
+A background thread polls each node at 1 Hz, collecting:
+
+- GPU utilisation (%)
+- Temperature (C)
+- System power draw (W)
+- CPU cluster usage (performance and efficiency cores)
+
+**Energy** is computed via trapezoidal integration of the power samples over each inference window (the wall-clock span of each benchmark request or concurrent batch). Average power is `total_joules / total_inference_seconds`.
+
+---
+
+## Output Format
+
+Results are written as JSON with three top-level keys:
+
+- **`runs`**: Array of per-request result objects, each containing:
+  - `elapsed_s`, `output_text_preview` (first 200 chars)
+  - `stats`: `{ prompt_tps, generation_tps, prompt_tokens, generation_tokens, peak_memory_usage }`
+  - Placement metadata: `model_id`, `placement_sharding`, `placement_instance_meta`, `placement_nodes`
+  - Run metadata: `pp_tokens`, `tg`, `repeat_index`, `concurrency`, `concurrent_index`
+  - `download_duration_s` (if model was freshly downloaded)
+- **`cluster`**: Cluster state snapshot at time of benchmarking.
+- **`system_metrics`**: Per-node time-series samples (GPU, power, temperature).
+
+---
+
+## Reproducing Results
+
+```bash
+cd bench && uv run python exo_bench.py \
+  --model "mlx-community/Qwen3.5-27B-4bit" \
+  --instance-meta jaccl \
+  --sharding tensor \
+  --min-nodes 2 --max-nodes 2 \
+  --pp 512 4096 --tg 128 \
+  --repeat 3 \
+  --warmup 1
+```
+
+Run --help for all the available flags.

bench/eval_tool_calls.py

@@ -3,11 +3,13 @@ from __future__ import annotations
 
 import argparse
 import contextlib
+import io
 import json
 import os
 import sys
 import time
 import tomllib
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Literal
@@ -17,6 +19,7 @@ from harness import (
     ExoClient,
     ExoHttpError,
     add_common_instance_args,
+    capture_cluster_snapshot,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
@@ -208,7 +211,7 @@ def _openai_build_request(
         "model": model,
         "messages": messages,
         "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     return "/v1/chat/completions", body
@@ -275,7 +278,7 @@ def _openai_build_followup(
         "model": model,
         "messages": followup_messages,
         "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     return "/v1/chat/completions", body
@@ -378,7 +381,7 @@ def _claude_build_request(
         "model": model,
         "messages": claude_messages,
         "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     if system_content is not None:
@@ -488,7 +491,7 @@ def _claude_build_followup(
         "model": model,
         "messages": claude_messages,
         "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     if system_content is not None:
@@ -912,6 +915,12 @@ Examples:
         default=1,
         help="Repeat each scenario N times (default: 1)",
     )
+    parser.add_argument(
+        "--concurrency",
+        type=int,
+        default=1,
+        help="Run up to N scenarios in parallel against the same instance (default: 1)",
+    )
     parser.add_argument(
         "--scenarios",
         nargs="*",
@@ -934,6 +943,13 @@ Examples:
     )
     args = parser.parse_args()
 
+    if args.concurrency < 1:
+        print(
+            f"--concurrency must be >= 1 (got {args.concurrency})",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
     all_scenarios = load_scenarios(SCENARIOS_PATH)
     if args.scenarios:
         scenarios = [s for s in all_scenarios if s.name in args.scenarios]
@@ -1006,44 +1022,75 @@ Examples:
         sys.exit(1)
 
     time.sleep(1)
+    cluster_snapshot = capture_cluster_snapshot(exo)
     all_results: list[ScenarioResult] = []
 
+    tasks: list[tuple[int, Scenario, ApiName]] = [
+        (run_idx, scenario, api_name)
+        for run_idx in range(args.repeat)
+        for scenario in scenarios
+        for api_name in api_names
+    ]
+
+    def _run_one(
+        http_client: httpx.Client,
+        task: tuple[int, Scenario, ApiName],
+    ) -> tuple[tuple[int, Scenario, ApiName], list[ScenarioResult], str]:
+        run_idx, scenario, api_name = task
+        buf = io.StringIO()
+        run_tag = f"[run {run_idx + 1}/{args.repeat}]" if args.repeat > 1 else ""
+        print(
+            f"\n  {run_tag}[{api_name:>9}] {scenario.name}: {scenario.description}",
+            file=buf,
+        )
+        scenario_results = run_scenario(
+            http_client,
+            args.host,
+            args.port,
+            full_model_id,
+            scenario,
+            api_name,
+            args.timeout,
+            args.verbose,
+        )
+        for r in scenario_results:
+            status = "PASS" if r.passed else "FAIL"
+            print(
+                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
+                file=buf,
+            )
+            for check_name, check_ok in r.checks.items():
+                mark = "+" if check_ok else "-"
+                print(f"      {mark} {check_name}", file=buf)
+            if r.error:
+                print(f"      ! {r.error}", file=buf)
+        return task, scenario_results, buf.getvalue()
+
     try:
         with httpx.Client() as http_client:
-            for run_idx in range(args.repeat):
-                if args.repeat > 1:
-                    print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
-
-                for scenario in scenarios:
-                    for api_name in api_names:
-                        print(
-                            f"\n  [{api_name:>9}] {scenario.name}: {scenario.description}",
-                            file=log,
-                        )
-
-                        scenario_results = run_scenario(
-                            http_client,
-                            args.host,
-                            args.port,
-                            full_model_id,
-                            scenario,
-                            api_name,
-                            args.timeout,
-                            args.verbose,
-                        )
+            if args.concurrency == 1:
+                current_run = -1
+                for task in tasks:
+                    run_idx = task[0]
+                    if args.repeat > 1 and run_idx != current_run:
+                        print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
+                        current_run = run_idx
+                    _, scenario_results, buffered = _run_one(http_client, task)
+                    all_results.extend(scenario_results)
+                    log.write(buffered)
+                    log.flush()
+            else:
+                print(
+                    f"Running {len(tasks)} tasks with concurrency={args.concurrency}",
+                    file=log,
+                )
+                with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
+                    futures = [pool.submit(_run_one, http_client, t) for t in tasks]
+                    for fut in as_completed(futures):
+                        _, scenario_results, buffered = fut.result()
                         all_results.extend(scenario_results)
-
-                        for r in scenario_results:
-                            status = "PASS" if r.passed else "FAIL"
-                            print(
-                                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
-                                file=log,
-                            )
-                            for check_name, check_ok in r.checks.items():
-                                mark = "+" if check_ok else "-"
-                                print(f"      {mark} {check_name}", file=log)
-                            if r.error:
-                                print(f"      ! {r.error}", file=log)
+                        log.write(buffered)
+                        log.flush()
     finally:
         try:
             exo.request_json("DELETE", f"/instance/{instance_id}")
@@ -1084,16 +1131,19 @@ Examples:
                 print(f"  - {r.name} [{r.api}/{r.phase}]: {r.error}", file=log)
 
     json_results = [result_to_dict(r) for r in all_results]
+    output: dict[str, Any] = {"results": json_results}
+    if cluster_snapshot:
+        output["cluster"] = cluster_snapshot
 
     if args.stdout:
-        print(json.dumps(json_results, indent=2))
+        print(json.dumps(output, indent=2))
     else:
         json_path = args.json_out
         parent = os.path.dirname(json_path)
         if parent:
             os.makedirs(parent, exist_ok=True)
         with open(json_path, "w") as f:
-            json.dump(json_results, f, indent=2)
+            json.dump(output, f, indent=2)
             f.write("\n")
         print(f"\nJSON results written to {json_path}", file=log)
 

bench/exo_bench.py

@@ -22,6 +22,7 @@ import contextlib
 import itertools
 import json
 import sys
+import threading
 import time
 from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -33,7 +34,9 @@ from harness import (
     ExoClient,
     ExoHttpError,
     add_common_instance_args,
+    capture_cluster_snapshot,
     instance_id_from_instance,
+    node_ids_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
     run_planning_phase,
@@ -131,6 +134,91 @@ def format_peak_memory(b: float) -> str:
     raise ValueError("You're using petabytes of memory. Something went wrong...")
 
 
+_SAMPLER_METRICS = ("gpuUsage", "temp", "sysPower", "pcpuUsage", "ecpuUsage")
+
+
+class SystemMetricsSampler:
+    def __init__(self, client: ExoClient, node_ids: list[str], interval_s: float = 1.0):
+        self._client = client
+        self._node_ids = node_ids
+        self._interval_s = interval_s
+        self._samples: dict[str, list[tuple[float, dict[str, float]]]] = {
+            nid: [] for nid in node_ids
+        }
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> None:
+        self._stop.clear()
+        self._thread = threading.Thread(target=self._poll_loop, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._stop.set()
+        if self._thread:
+            self._thread.join(timeout=5)
+
+    def _poll_loop(self) -> None:
+        while not self._stop.is_set():
+            t = time.monotonic()
+            for nid in self._node_ids:
+                try:
+                    data = self._client.get_node_system(nid)
+                    if data:
+                        self._samples[nid].append(
+                            (t, {k: data.get(k, 0.0) for k in _SAMPLER_METRICS})
+                        )
+                except Exception:
+                    pass
+            self._stop.wait(self._interval_s)
+
+    def energy_between(self, t0: float, t1: float) -> float:
+        total_joules = 0.0
+        for _nid, samples in self._samples.items():
+            window = [(t, s["sysPower"]) for t, s in samples if t0 <= t <= t1]
+            if len(window) >= 2:
+                for i in range(1, len(window)):
+                    dt = window[i][0] - window[i - 1][0]
+                    avg_power = (window[i][1] + window[i - 1][1]) / 2
+                    total_joules += avg_power * dt
+            elif len(window) == 1:
+                total_joules += window[0][1] * (t1 - t0)
+        return total_joules
+
+    def summarize(self) -> dict[str, dict[str, dict[str, float]]]:
+        result: dict[str, dict[str, dict[str, float]]] = {}
+        for nid, samples in self._samples.items():
+            if not samples:
+                continue
+            metrics: dict[str, dict[str, float]] = {}
+            for key in _SAMPLER_METRICS:
+                values = [s[key] for t, s in samples]
+                metrics[key] = {
+                    "min": round(min(values), 2),
+                    "max": round(max(values), 2),
+                    "mean": round(mean(values), 2),
+                    "samples": len(values),
+                }
+            result[nid] = metrics
+        return result
+
+    def print_summary(self, placement_label: str) -> None:
+        summary = self.summarize()
+        if not summary:
+            return
+        logger.info(f"--- System Metrics ({placement_label}) ---")
+        for nid, metrics in summary.items():
+            gpu = metrics.get("gpuUsage", {})
+            temp = metrics.get("temp", {})
+            power = metrics.get("sysPower", {})
+            logger.info(
+                f"  {nid}:  "
+                f"GPU {gpu.get('mean', 0) * 100:.0f}% avg ({gpu.get('min', 0) * 100:.0f}–{gpu.get('max', 0) * 100:.0f}%)  |  "
+                f"{temp.get('mean', 0):.1f}°C avg  |  "
+                f"{power.get('mean', 0):.1f}W avg"
+            )
+
+
 def parse_int_list(values: list[str]) -> list[int]:
     items: list[int] = []
     for v in values:
@@ -142,7 +230,13 @@ def parse_int_list(values: list[str]) -> list[int]:
 
 
 def run_one_completion(
-    client: ExoClient, model_id: str, pp_hint: int, tg: int, prompt_sizer: PromptSizer
+    client: ExoClient,
+    model_id: str,
+    pp_hint: int,
+    tg: int,
+    prompt_sizer: PromptSizer,
+    *,
+    use_prefix_cache: bool = False,
 ) -> tuple[dict[str, Any], int]:
     content, pp_tokens = prompt_sizer.build(pp_hint)
     payload: dict[str, Any] = {
@@ -150,6 +244,8 @@ def run_one_completion(
         "messages": [{"role": "user", "content": content}],
         "stream": False,
         "max_tokens": tg,
+        "logprobs": False,
+        "use_prefix_cache": use_prefix_cache,
     }
 
     t0 = time.perf_counter()
@@ -279,6 +375,22 @@ def main() -> int:
         action="store_true",
         help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
     )
+    ap.add_argument(
+        "--no-system-metrics",
+        action="store_true",
+        help="Disable GPU utilization, temperature, and power collection during inference.",
+    )
+    ap.add_argument(
+        "--metrics-interval",
+        type=float,
+        default=1.0,
+        help="System metrics polling interval in seconds (default: 1.0).",
+    )
+    ap.add_argument(
+        "--use-prefix-cache",
+        action="store_true",
+        help="Enable KV prefix cache during bench (default: disabled for cold-cache measurements).",
+    )
     args = ap.parse_args()
 
     pp_list = parse_int_list(args.pp)
@@ -294,6 +406,15 @@ def main() -> int:
         logger.error("--concurrency values must be >= 1")
         return 2
 
+    if args.use_prefix_cache:
+        logger.warning(
+            "--use-prefix-cache: prompt TPS will be approximate. See METHODOLOGY.md for details."
+        )
+        if pp_list != sorted(pp_list):
+            logger.warning(
+                "--pp values are not in ascending order: prompt TPS will be less accurate. Use ascending --pp for best results."
+            )
+
     # Log pairing mode
     use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
     if use_combinations:
@@ -364,7 +485,9 @@ def main() -> int:
     else:
         logger.info("Download: model already cached")
 
+    cluster_snapshot = capture_cluster_snapshot(client)
     all_rows: list[dict[str, Any]] = []
+    all_system_metrics: dict[str, dict[str, dict[str, float]]] = {}
 
     for preview in selected:
         instance = preview["instance"]
@@ -390,10 +513,25 @@ def main() -> int:
 
         time.sleep(1)
 
+        sampler: SystemMetricsSampler | None = None
+        if not args.no_system_metrics:
+            nids = node_ids_from_instance(instance)
+            sampler = SystemMetricsSampler(
+                ExoClient(args.host, args.port, timeout_s=30),
+                nids,
+                interval_s=args.metrics_interval,
+            )
+            sampler.start()
+
         try:
             for i in range(args.warmup):
                 run_one_completion(
-                    client, full_model_id, pp_list[0], tg_list[0], prompt_sizer
+                    client,
+                    full_model_id,
+                    pp_list[0],
+                    tg_list[0],
+                    prompt_sizer,
+                    use_prefix_cache=args.use_prefix_cache,
                 )
                 logger.debug(f"  warmup {i + 1}/{args.warmup} done")
 
@@ -408,15 +546,23 @@ def main() -> int:
                 for concurrency in concurrency_list:
                     logger.info(f"--- pp={pp} tg={tg} concurrency={concurrency} ---")
                     runs: list[dict[str, Any]] = []
+                    inference_windows: list[tuple[float, float]] = []
                     for r in range(args.repeat):
                         time.sleep(3)
 
                         if concurrency <= 1:
                             # Sequential: single request
                             try:
+                                inf_t0 = time.monotonic()
                                 row, actual_pp_tokens = run_one_completion(
-                                    client, full_model_id, pp, tg, prompt_sizer
+                                    client,
+                                    full_model_id,
+                                    pp,
+                                    tg,
+                                    prompt_sizer,
+                                    use_prefix_cache=args.use_prefix_cache,
                                 )
+                                inference_windows.append((inf_t0, time.monotonic()))
                             except Exception as e:
                                 logger.error(e)
                                 continue
@@ -443,20 +589,54 @@ def main() -> int:
                             all_rows.append(row)
                         else:
                             # Concurrent: fire N requests in parallel
-                            # Each thread gets its own ExoClient (separate HTTP connection)
+                            # Pre-build prompt once, barrier ensures simultaneous dispatch
+                            content, actual_pp = prompt_sizer.build(pp)
+                            pre_built_payload: dict[str, Any] = {
+                                "model": full_model_id,
+                                "messages": [{"role": "user", "content": content}],
+                                "stream": False,
+                                "max_tokens": tg,
+                                "logprobs": False,
+                                "use_prefix_cache": args.use_prefix_cache,
+                            }
+                            barrier = threading.Barrier(concurrency)
+                            batch_start = threading.Event()
+                            batch_t0: float = 0.0
                             batch_results: list[tuple[dict[str, Any], int]] = []
                             batch_errors = 0
 
                             def _run_concurrent(
-                                idx: int, *, _pp: int = pp, _tg: int = tg
+                                idx: int,
+                                _barrier: threading.Barrier = barrier,
+                                _batch_start: threading.Event = batch_start,
+                                _payload: dict[str, Any] = pre_built_payload,
+                                _actual_pp: int = actual_pp,
                             ) -> tuple[dict[str, Any], int]:
+                                nonlocal batch_t0
                                 c = ExoClient(
                                     args.host, args.port, timeout_s=args.timeout
                                 )
-                                return run_one_completion(
-                                    c, full_model_id, _pp, _tg, prompt_sizer
+                                if _barrier.wait() == 0:
+                                    batch_t0 = time.perf_counter()
+                                    _batch_start.set()
+                                else:
+                                    _batch_start.wait()
+                                t0 = batch_t0
+                                out = c.post_bench_chat_completions(_payload)
+                                elapsed = time.perf_counter() - t0
+                                stats = out.get("generation_stats")
+                                choices = out.get("choices") or [{}]
+                                message = (
+                                    choices[0].get("message", {}) if choices else {}
                                 )
+                                text = message.get("content") or ""
+                                return {
+                                    "elapsed_s": elapsed,
+                                    "output_text_preview": text[:200],
+                                    "stats": stats,
+                                }, _actual_pp
 
+                            inf_t0 = time.monotonic()
                             with ThreadPoolExecutor(max_workers=concurrency) as pool:
                                 futures = {
                                     pool.submit(_run_concurrent, i): i
@@ -468,6 +648,12 @@ def main() -> int:
                                     except Exception as e:
                                         logger.error(f"Concurrent request failed: {e}")
                                         batch_errors += 1
+                            batch_wall_s = (
+                                max(x["elapsed_s"] for x, _ in batch_results)
+                                if batch_results
+                                else time.perf_counter() - batch_t0
+                            )
+                            inference_windows.append((inf_t0, time.monotonic()))
 
                             for idx, (row, actual_pp_tokens) in enumerate(
                                 batch_results
@@ -501,36 +687,57 @@ def main() -> int:
                                     for x, _ in batch_results
                                     if x["stats"]["generation_tps"] > 0
                                 ]
-                                agg_gen_tps = (
-                                    mean(valid_gen_tps) if valid_gen_tps else 0.0
+                                per_req_tps = (
+                                    max(valid_gen_tps) if valid_gen_tps else 0.0
                                 )
-                                gen_tps = agg_gen_tps / concurrency
+                                agg_gen_tps = per_req_tps * concurrency
                                 logger.info(
                                     f"[concurrent {concurrency}x]  "
                                     f"agg_gen_tps={agg_gen_tps:.2f}  "
-                                    f"gen_tps={gen_tps:.2f}  "
+                                    f"per_req_tps={per_req_tps:.2f}  "
+                                    f"wall_s={batch_wall_s:.2f}  "
                                     f"errors={batch_errors}"
                                 )
 
                     if runs:
                         prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-                        gen_tps = mean(
-                            x["stats"]["generation_tps"] / x["concurrency"]
+                        valid_gen = [
+                            x["stats"]["generation_tps"]
                             for x in runs
-                        )
+                            if x["stats"]["generation_tps"] > 0
+                        ]
+                        per_req_tps = max(valid_gen) if valid_gen else 0.0
+                        gen_tps = per_req_tps * concurrency
                         ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
                         gtok = mean(x["stats"]["generation_tokens"] for x in runs)
                         peak = mean(
                             x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
                         )
 
-                        logger.info(
+                        summary = (
                             f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
                             f"prompt_tokens={ptok} gen_tokens={gtok}    "
-                            f"peak_memory={format_peak_memory(peak)}\n"
+                            f"peak_memory={format_peak_memory(peak)}"
                         )
+                        if sampler and inference_windows:
+                            joules = sum(
+                                sampler.energy_between(t0, t1)
+                                for t0, t1 in inference_windows
+                            )
+                            inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
+                            avg_watts = joules / inf_seconds if inf_seconds > 0 else 0
+                            summary += f"    energy={joules:.1f}J ({avg_watts:.1f}W avg over {inf_seconds:.1f}s inference)"
+                        logger.info(f"{summary}\n")
                     time.sleep(2)
         finally:
+            if sampler:
+                sampler.stop()
+                placement_label = f"{sharding}/{instance_meta}/{n_nodes} nodes"
+                sampler.print_summary(placement_label)
+                placement_metrics = sampler.summarize()
+                if placement_metrics:
+                    all_system_metrics.update(placement_metrics)
+
             try:
                 client.request_json("DELETE", f"/instance/{instance_id}")
             except ExoHttpError as e:
@@ -541,11 +748,17 @@ def main() -> int:
 
             time.sleep(5)
 
+    output: dict[str, Any] = {"runs": all_rows}
+    if cluster_snapshot:
+        output["cluster"] = cluster_snapshot
+    if all_system_metrics:
+        output["system_metrics"] = all_system_metrics
+
     if args.stdout:
-        json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False)
+        json.dump(output, sys.stdout, indent=2, ensure_ascii=False)
     elif args.json_out:
         with open(args.json_out, "w", encoding="utf-8") as f:
-            json.dump(all_rows, f, indent=2, ensure_ascii=False)
+            json.dump(output, f, indent=2, ensure_ascii=False)
         logger.debug(f"\nWrote results JSON: {args.json_out}")
 
     return 0

bench/exo_eval.py

@@ -46,6 +46,7 @@ from harness import (
     ExoClient,
     ExoHttpError,
     add_common_instance_args,
+    capture_cluster_snapshot,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
@@ -1027,16 +1028,18 @@ def save_results(
     concurrency: int,
     results: list[QuestionResult],
     scores: dict[str, Any],
+    cluster: dict[str, Any] | None = None,
 ) -> Path:
     out_dir = Path(results_dir) / model.replace("/", "_") / benchmark_name
     out_dir.mkdir(parents=True, exist_ok=True)
     ts = time.strftime("%Y%m%d_%H%M%S")
     path = out_dir / f"c{concurrency}_{ts}.json"
 
-    data = {
+    data: dict[str, Any] = {
         "benchmark": benchmark_name,
         "model": model,
         "concurrency": concurrency,
+        **({"cluster": cluster} if cluster else {}),
         "scores": scores,
         "results": [
             {
@@ -1231,8 +1234,10 @@ def main() -> int:
                 client.request_json("DELETE", f"/instance/{instance_id}")
             return 1
         time.sleep(1)
+        cluster_snapshot = capture_cluster_snapshot(client)
     else:
         full_model_id = args.model
+        cluster_snapshot = None
 
     # Auto-detect reasoning from model config
     model_config = load_model_config(full_model_id)
@@ -1328,6 +1333,7 @@ def main() -> int:
                             c,
                             results,
                             scores,
+                            cluster=cluster_snapshot,
                         )
                         results_by_c[c] = results
                 if len(results_by_c) >= 2:
@@ -1358,6 +1364,7 @@ def main() -> int:
                         args.num_concurrent,
                         results,
                         scores,
+                        cluster=cluster_snapshot,
                     )
     finally:
         if instance_id is not None:

bench/harness.py

@@ -69,6 +69,35 @@ class ExoClient:
     def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
         return self.request_json("POST", "/bench/chat/completions", body=payload)
 
+    def get_state_path(self, path: str) -> Any:
+        try:
+            return self.request_json("GET", f"/state/{path}")
+        except ExoHttpError as e:
+            if e.status == 404:
+                return None
+            raise
+
+    def get_instance(self, instance_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"instances/{instance_id}")
+
+    def get_runner(self, runner_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"runners/{runner_id}")
+
+    def get_node_downloads(self, node_id: str) -> list[dict[str, Any]] | None:
+        return self.get_state_path(f"downloads/{node_id}")
+
+    def get_node_disk(self, node_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"nodeDisk/{node_id}")
+
+    def get_node_system(self, node_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"nodeSystem/{node_id}")
+
+    def get_node_identities(self) -> dict[str, Any] | None:
+        return self.get_state_path("nodeIdentities")
+
+    def get_topology(self) -> dict[str, Any] | None:
+        return self.get_state_path("topology")
+
 
 def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
     if len(instance) != 1:
@@ -97,6 +126,11 @@ def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]:
     return list(runner_to_shard.keys())
 
 
+def node_ids_from_instance(instance: dict[str, Any]) -> list[str]:
+    inner = unwrap_instance(instance)
+    return list(inner["shardAssignments"]["nodeToRunner"].keys())
+
+
 def runner_ready(runner: dict[str, Any]) -> bool:
     return "RunnerReady" in runner
 
@@ -116,13 +150,12 @@ def wait_for_instance_ready(
 ) -> None:
     start_time = time.time()
     instance_existed = False
+    last_loaded: dict[str, int] = {}
     while time.time() - start_time < timeout:
-        state = client.request_json("GET", "/state")
-        instances = state.get("instances", {})
+        instance = client.get_instance(instance_id)
 
-        if instance_id not in instances:
+        if instance is None:
             if instance_existed:
-                # Instance was deleted after being created - likely due to runner failure
                 raise RuntimeError(
                     f"Instance {instance_id} was deleted (runner may have failed)"
                 )
@@ -130,18 +163,25 @@ def wait_for_instance_ready(
             continue
 
         instance_existed = True
-        instance = instances[instance_id]
-        runner_ids = runner_ids_from_instance(instance)
-        runners = state.get("runners", {})
+        rids = runner_ids_from_instance(instance)
 
-        # Check for failed runners first
-        for rid in runner_ids:
-            runner = runners.get(rid, {})
+        all_ready = True
+        for rid in rids:
+            runner = client.get_runner(rid) or {}
             if runner_failed(runner):
                 error_msg = get_runner_failed_message(runner) or "Unknown error"
                 raise RuntimeError(f"Runner {rid} failed: {error_msg}")
+            if "RunnerLoading" in runner:
+                loading = runner["RunnerLoading"]
+                loaded = loading.get("layersLoaded", 0)
+                total = loading.get("totalLayers", 0)
+                if total > 0 and last_loaded.get(rid) != loaded:
+                    last_loaded[rid] = loaded
+                    logger.debug(f"Runner {rid}: loading layers {loaded}/{total}")
+            if not runner_ready(runner):
+                all_ready = False
 
-        if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
+        if all_ready:
             return
 
         time.sleep(0.1)
@@ -165,6 +205,23 @@ def wait_for_instance_gone(
     raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}")
 
 
+def capture_cluster_snapshot(client: ExoClient) -> dict[str, Any]:
+    snapshot: dict[str, Any] = {}
+    identities = client.get_node_identities()
+    if identities:
+        snapshot["nodeIdentities"] = identities
+    topology = client.get_topology()
+    if topology:
+        snapshot["topology"] = topology
+    node_memory = client.get_state_path("nodeMemory")
+    if node_memory:
+        snapshot["nodeMemory"] = node_memory
+    node_system = client.get_state_path("nodeSystem")
+    if node_system:
+        snapshot["nodeSystem"] = node_system
+    return snapshot
+
+
 def resolve_model_short_id(
     client: ExoClient, model_arg: str, *, force_download: bool = False
 ) -> tuple[str, str]:
@@ -326,16 +383,11 @@ def run_planning_phase(
     node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
     runner_to_shard = inner["shardAssignments"]["runnerToShard"]
 
-    state = client.request_json("GET", "/state")
-    downloads = state.get("downloads", {})
-    node_disk = state.get("nodeDisk", {})
-
     needs_download = False
 
     for node_id in node_ids:
-        node_downloads = downloads.get(node_id, [])
+        node_downloads = client.get_node_downloads(node_id) or []
 
-        # Check if model already downloaded on this node
         already_downloaded = any(
             "DownloadCompleted" in p
             and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
@@ -349,8 +401,7 @@ def run_planning_phase(
 
         needs_download = True
 
-        # Wait for disk info if settle_deadline is set
-        disk_info = node_disk.get(node_id, {})
+        disk_info = client.get_node_disk(node_id) or {}
         backoff = _SETTLE_INITIAL_BACKOFF_S
         while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
             remaining = settle_deadline - time.monotonic()
@@ -359,9 +410,7 @@ def run_planning_phase(
             )
             time.sleep(min(backoff, remaining))
             backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
-            state = client.request_json("GET", "/state")
-            node_disk = state.get("nodeDisk", {})
-            disk_info = node_disk.get(node_id, {})
+            disk_info = client.get_node_disk(node_id) or {}
 
         if not disk_info:
             logger.warning(f"No disk info for {node_id}, skipping space check")
@@ -377,7 +426,6 @@ def run_planning_phase(
                 f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
             )
 
-        # Delete from smallest to largest (skip read-only models from EXO_MODELS_PATH)
         completed = [
             (
                 unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
@@ -417,21 +465,20 @@ def run_planning_phase(
     # Wait for downloads
     start = time.time()
     while time.time() - start < timeout:
-        state = client.request_json("GET", "/state")
-        downloads = state.get("downloads", {})
         all_done = True
         for node_id in node_ids:
+            node_downloads = client.get_node_downloads(node_id) or []
             done = any(
                 "DownloadCompleted" in p
                 and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
                     "modelCard"
                 ]["modelId"]
                 == full_model_id
-                for p in downloads.get(node_id, [])
+                for p in node_downloads
             )
             failed = [
                 p["DownloadFailed"]["errorMessage"]
-                for p in downloads.get(node_id, [])
+                for p in node_downloads
                 if "DownloadFailed" in p
                 and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
                     "modelId"
@@ -442,6 +489,27 @@ def run_planning_phase(
                 raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
             if not done:
                 all_done = False
+                ongoing = [
+                    p
+                    for p in node_downloads
+                    if "DownloadOngoing" in p
+                    and unwrap_instance(p["DownloadOngoing"]["shardMetadata"])[
+                        "modelCard"
+                    ]["modelId"]
+                    == full_model_id
+                ]
+                if ongoing:
+                    prog = ongoing[0]["DownloadOngoing"]["downloadProgress"]
+                    speed_mb = prog.get("speed", 0) / (1024 * 1024)
+                    eta_s = prog.get("etaMs", 0) / 1000
+                    dl_bytes = prog.get("downloaded", {}).get("inBytes", 0)
+                    total_bytes = prog.get("total", {}).get("inBytes", 0)
+                    pct = (dl_bytes / total_bytes * 100) if total_bytes else 0
+                    logger.info(
+                        f"Downloading on {node_id}: {pct:.1f}% @ {speed_mb:.1f} MB/s, "
+                        f"ETA {eta_s:.0f}s "
+                        f"({prog.get('completedFiles', 0)}/{prog.get('totalFiles', 0)} files)"
+                    )
         if all_done:
             if download_t0 is not None:
                 return time.perf_counter() - download_t0
@@ -496,7 +564,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
     ap.add_argument(
         "--settle-timeout",
         type=float,
-        default=0,
+        default=60.0,
         help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
     )
     ap.add_argument(

dashboard/package-lock.json

@@ -11,7 +11,8 @@
 				"highlight.js": "^11.11.1",
 				"katex": "^0.16.27",
 				"marked": "^17.0.1",
-				"mode-watcher": "^1.1.0"
+				"mode-watcher": "^1.1.0",
+				"pdfjs-dist": "^5.6.205"
 			},
 			"devDependencies": {
 				"@sveltejs/adapter-static": "^3.0.10",
@@ -518,6 +519,256 @@
 				"@jridgewell/sourcemap-codec": "^1.4.14"
 			}
 		},
+		"node_modules/@napi-rs/canvas": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas/-/canvas-0.1.97.tgz",
+			"integrity": "sha512-8cFniXvrIEnVwuNSRCW9wirRZbHvrD3JVujdS2P5n5xiJZNZMOZcfOvJ1pb66c7jXMKHHglJEDVJGbm8XWFcXQ==",
+			"license": "MIT",
+			"optional": true,
+			"workspaces": [
+				"e2e/*"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			},
+			"optionalDependencies": {
+				"@napi-rs/canvas-android-arm64": "0.1.97",
+				"@napi-rs/canvas-darwin-arm64": "0.1.97",
+				"@napi-rs/canvas-darwin-x64": "0.1.97",
+				"@napi-rs/canvas-linux-arm-gnueabihf": "0.1.97",
+				"@napi-rs/canvas-linux-arm64-gnu": "0.1.97",
+				"@napi-rs/canvas-linux-arm64-musl": "0.1.97",
+				"@napi-rs/canvas-linux-riscv64-gnu": "0.1.97",
+				"@napi-rs/canvas-linux-x64-gnu": "0.1.97",
+				"@napi-rs/canvas-linux-x64-musl": "0.1.97",
+				"@napi-rs/canvas-win32-arm64-msvc": "0.1.97",
+				"@napi-rs/canvas-win32-x64-msvc": "0.1.97"
+			}
+		},
+		"node_modules/@napi-rs/canvas-android-arm64": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-android-arm64/-/canvas-android-arm64-0.1.97.tgz",
+			"integrity": "sha512-V1c/WVw+NzH8vk7ZK/O8/nyBSCQimU8sfMsB/9qeSvdkGKNU7+mxy/bIF0gTgeBFmHpj30S4E9WHMSrxXGQuVQ==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"android"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-darwin-arm64": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-arm64/-/canvas-darwin-arm64-0.1.97.tgz",
+			"integrity": "sha512-ok+SCEF4YejcxuJ9Rm+WWunHHpf2HmiPxfz6z1a/NFQECGXtsY7A4B8XocK1LmT1D7P174MzwPF9Wy3AUAwEPw==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"darwin"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-darwin-x64": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-x64/-/canvas-darwin-x64-0.1.97.tgz",
+			"integrity": "sha512-PUP6e6/UGlclUvAQNnuXCcnkpdUou6VYZfQOQxExLp86epOylmiwLkqXIvpFmjoTEDmPmXrI+coL/9EFU1gKPA==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"darwin"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-arm-gnueabihf": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm-gnueabihf/-/canvas-linux-arm-gnueabihf-0.1.97.tgz",
+			"integrity": "sha512-XyXH2L/cic8eTNtbrXCcvqHtMX/nEOxN18+7rMrAM2XtLYC/EB5s0wnO1FsLMWmK+04ZSLN9FBGipo7kpIkcOw==",
+			"cpu": [
+				"arm"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-arm64-gnu": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-gnu/-/canvas-linux-arm64-gnu-0.1.97.tgz",
+			"integrity": "sha512-Kuq/M3djq0K8ktgz6nPlK7Ne5d4uWeDxPpyKWOjWDK2RIOhHVtLtyLiJw2fuldw7Vn4mhw05EZXCEr4Q76rs9w==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-arm64-musl": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-musl/-/canvas-linux-arm64-musl-0.1.97.tgz",
+			"integrity": "sha512-kKmSkQVnWeqg7qdsiXvYxKhAFuHz3tkBjW/zyQv5YKUPhotpaVhpBGv5LqCngzyuRV85SXoe+OFj+Tv0a0QXkQ==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-riscv64-gnu": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-riscv64-gnu/-/canvas-linux-riscv64-gnu-0.1.97.tgz",
+			"integrity": "sha512-Jc7I3A51jnEOIAXeLsN/M/+Z28LUeakcsXs07FLq9prXc0eYOtVwsDEv913Gr+06IRo34gJJVgT0TXvmz+N2VA==",
+			"cpu": [
+				"riscv64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-x64-gnu": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-gnu/-/canvas-linux-x64-gnu-0.1.97.tgz",
+			"integrity": "sha512-iDUBe7AilfuBSRbSa8/IGX38Mf+iCSBqoVKLSQ5XaY2JLOaqz1TVyPFEyIck7wT6mRQhQt5sN6ogfjIDfi74tg==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-x64-musl": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-musl/-/canvas-linux-x64-musl-0.1.97.tgz",
+			"integrity": "sha512-AKLFd/v0Z5fvgqBDqhvqtAdx+fHMJ5t9JcUNKq4FIZ5WH+iegGm8HPdj00NFlCSnm83Fp3Ln8I2f7uq1aIiWaA==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-win32-arm64-msvc": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-arm64-msvc/-/canvas-win32-arm64-msvc-0.1.97.tgz",
+			"integrity": "sha512-u883Yr6A6fO7Vpsy9YE4FVCIxzzo5sO+7pIUjjoDLjS3vQaNMkVzx5bdIpEL+ob+gU88WDK4VcxYMZ6nmnoX9A==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"win32"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-win32-x64-msvc": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-x64-msvc/-/canvas-win32-x64-msvc-0.1.97.tgz",
+			"integrity": "sha512-sWtD2EE3fV0IzN+iiQUqr/Q1SwqWhs2O1FKItFlxtdDkikpEj5g7DKQpY3x55H/MAOnL8iomnlk3mcEeGiUMoQ==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"win32"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
 		"node_modules/@polka/url": {
 			"version": "1.0.0-next.29",
 			"resolved": "https://registry.npmjs.org/@polka/url/-/url-1.0.0-next.29.tgz",
@@ -2635,6 +2886,26 @@
 				"node": "^10 || ^12 || ^13.7 || ^14 || >=15.0.1"
 			}
 		},
+		"node_modules/node-readable-to-web-readable-stream": {
+			"version": "0.4.2",
+			"resolved": "https://registry.npmjs.org/node-readable-to-web-readable-stream/-/node-readable-to-web-readable-stream-0.4.2.tgz",
+			"integrity": "sha512-/cMZNI34v//jUTrI+UIo4ieHAB5EZRY/+7OmXZgBxaWBMcW2tGdceIw06RFxWxrKZ5Jp3sI2i5TsRo+CBhtVLQ==",
+			"license": "MIT",
+			"optional": true
+		},
+		"node_modules/pdfjs-dist": {
+			"version": "5.6.205",
+			"resolved": "https://registry.npmjs.org/pdfjs-dist/-/pdfjs-dist-5.6.205.tgz",
+			"integrity": "sha512-tlUj+2IDa7G1SbvBNN74UHRLJybZDWYom+k6p5KIZl7huBvsA4APi6mKL+zCxd3tLjN5hOOEE9Tv7VdzO88pfg==",
+			"license": "Apache-2.0",
+			"engines": {
+				"node": ">=20.19.0 || >=22.13.0 || >=24"
+			},
+			"optionalDependencies": {
+				"@napi-rs/canvas": "^0.1.96",
+				"node-readable-to-web-readable-stream": "^0.4.2"
+			}
+		},
 		"node_modules/picocolors": {
 			"version": "1.1.1",
 			"resolved": "https://registry.npmjs.org/picocolors/-/picocolors-1.1.1.tgz",

dashboard/package.json

@@ -31,6 +31,7 @@
 		"highlight.js": "^11.11.1",
 		"katex": "^0.16.27",
 		"marked": "^17.0.1",
-		"mode-watcher": "^1.1.0"
+		"mode-watcher": "^1.1.0",
+		"pdfjs-dist": "^5.6.205"
 	}
 }

dashboard/src/lib/components/ChatMessages.svelte

@@ -139,6 +139,8 @@
         return "🖼";
       case "text":
         return "📄";
+      case "pdf":
+        return "📑";
       default:
         return "📎";
     }

dashboard/src/lib/components/FamilyLogos.svelte

@@ -88,6 +88,18 @@
       d="M22.012 0h1.032v.927H24v.968h-.956V3.78h-1.032V1.896h-1.878v-.97h1.878V0zM2.6 12.371V1.87h.969v10.502h-.97zm10.423.66h10.95v.918h-6.208v9.579h-4.742V13.03zM5.629 3.333v12.356H0v4.51h10.386V8L20.859 8l-.003-4.668-15.227.001z"
     />
   </svg>
+{:else if family === "nemotron"}
+  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
+    <path
+      d="M8.948 8.798v-1.43a6.7 6.7 0 0 1 .424-.018c3.922-.124 6.493 3.374 6.493 3.374s-2.774 3.851-5.75 3.851c-.398 0-.787-.062-1.158-.185v-4.346c1.528.185 1.837.857 2.747 2.385l2.04-1.714s-1.492-1.952-4-1.952a6.016 6.016 0 0 0-.796.035m0-4.735v2.138l.424-.027c5.45-.185 9.01 4.47 9.01 4.47s-4.08 4.964-8.33 4.964c-.37 0-.733-.035-1.095-.097v1.325c.3.035.61.062.91.062 3.957 0 6.82-2.023 9.593-4.408.459.371 2.34 1.263 2.73 1.652-2.633 2.208-8.772 3.984-12.253 3.984-.335 0-.653-.018-.971-.053v1.864H24V4.063zm0 10.326v1.131c-3.657-.654-4.673-4.46-4.673-4.46s1.758-1.944 4.673-2.262v1.237H8.94c-1.528-.186-2.73 1.245-2.73 1.245s.68 2.412 2.739 3.11M2.456 10.9s2.164-3.197 6.5-3.533V6.201C4.153 6.59 0 10.653 0 10.653s2.35 6.802 8.948 7.42v-1.237c-4.84-.6-6.492-5.936-6.492-5.936z"
+    />
+  </svg>
+{:else if family === "gemma"}
+  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
+    <path
+      d="M12.48 10.92v3.28h7.84c-.24 1.84-.853 3.187-1.787 4.133-1.147 1.147-2.933 2.4-6.053 2.4-4.827 0-8.6-3.893-8.6-8.72s3.773-8.72 8.6-8.72c2.6 0 4.507 1.027 5.907 2.347l2.307-2.307C18.747 1.44 16.133 0 12.48 0 5.867 0 .307 5.387.307 12s5.56 12 12.173 12c3.573 0 6.267-1.173 8.373-3.36 2.16-2.16 2.84-5.213 2.84-7.667 0-.76-.053-1.467-.173-2.053H12.48z"
+    />
+  </svg>
 {:else}
   <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
     <path

dashboard/src/lib/components/FamilySidebar.svelte

@@ -31,6 +31,8 @@
     kimi: "Kimi",
     flux: "FLUX",
     "qwen-image": "Qwen Img",
+    nemotron: "NVIDIA",
+    gemma: "Google",
   };
 
   function getFamilyName(family: string): string {
@@ -41,31 +43,20 @@
 </script>
 
 <div
-  class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[72px] sm:min-w-[64px] overflow-y-auto scrollbar-hide"
+  class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[80px] sm:min-w-[72px] overflow-y-auto scrollbar-hide"
 >
   <!-- All models (no filter) -->
   <button
     type="button"
     onclick={() => onSelect(null)}
-    class="group flex flex-col items-center justify-center p-2 sm:p-2 rounded transition-all duration-200 cursor-pointer min-h-[44px] sm:min-h-0 {selectedFamily ===
+    class="group flex items-center justify-center px-3 py-2.5 rounded transition-all duration-200 cursor-pointer min-h-[44px] sm:min-h-0 {selectedFamily ===
     null
       ? 'bg-exo-yellow/20 border-l-2 border-exo-yellow'
       : 'hover:bg-white/5 border-l-2 border-transparent'}"
     title="All models"
   >
-    <svg
-      class="w-5 h-5 {selectedFamily === null
-        ? 'text-exo-yellow'
-        : 'text-white/50 group-hover:text-white/70'}"
-      viewBox="0 0 24 24"
-      fill="currentColor"
-    >
-      <path
-        d="M4 8h4V4H4v4zm6 12h4v-4h-4v4zm-6 0h4v-4H4v4zm0-6h4v-4H4v4zm6 0h4v-4h-4v4zm6-10v4h4V4h-4zm-6 4h4V4h-4v4zm6 6h4v-4h-4v4zm0 6h4v-4h-4v4z"
-      />
-    </svg>
     <span
-      class="text-[9px] font-mono mt-0.5 {selectedFamily === null
+      class="text-[12px] font-mono font-medium {selectedFamily === null
         ? 'text-exo-yellow'
         : 'text-white/40 group-hover:text-white/60'}">All</span
     >
@@ -89,7 +80,7 @@
           : "text-white/50 group-hover:text-amber-400/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 {selectedFamily === 'favorites'
+        class="text-[11px] font-mono mt-0.5 {selectedFamily === 'favorites'
           ? 'text-amber-400'
           : 'text-white/40 group-hover:text-white/60'}">Faves</span
       >
@@ -114,7 +105,7 @@
           : "text-white/50 group-hover:text-white/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 {selectedFamily === 'recents'
+        class="text-[11px] font-mono mt-0.5 {selectedFamily === 'recents'
           ? 'text-exo-yellow'
           : 'text-white/40 group-hover:text-white/60'}">Recent</span
       >
@@ -138,7 +129,7 @@
         : "text-white/50 group-hover:text-orange-400/70"}
     />
     <span
-      class="text-[9px] font-mono mt-0.5 {selectedFamily === 'huggingface'
+      class="text-[11px] font-mono mt-0.5 {selectedFamily === 'huggingface'
         ? 'text-orange-400'
         : 'text-white/40 group-hover:text-white/60'}">Hub</span
     >
@@ -164,7 +155,7 @@
           : "text-white/50 group-hover:text-white/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 truncate max-w-full {selectedFamily ===
+        class="text-[11px] font-mono mt-0.5 truncate max-w-full {selectedFamily ===
         family
           ? 'text-exo-yellow'
           : 'text-white/40 group-hover:text-white/60'}"

dashboard/src/lib/components/HeaderNav.svelte

@@ -1,21 +1,38 @@
 <script lang="ts">
   import { browser } from "$app/environment";
 
-  export let showHome = true;
-  export let onHome: (() => void) | null = null;
-  export let showSidebarToggle = false;
-  export let sidebarVisible = true;
-  export let onToggleSidebar: (() => void) | null = null;
-  export let showMobileMenuToggle = false;
-  export let mobileMenuOpen = false;
-  export let onToggleMobileMenu: (() => void) | null = null;
-  export let showMobileRightToggle = false;
-  export let mobileRightOpen = false;
-  export let onToggleMobileRight: (() => void) | null = null;
-  export let downloadProgress: {
-    count: number;
-    percentage: number;
-  } | null = null;
+  interface Props {
+    showHome?: boolean;
+    onHome?: (() => void) | null;
+    showSidebarToggle?: boolean;
+    sidebarVisible?: boolean;
+    onToggleSidebar?: (() => void) | null;
+    showMobileMenuToggle?: boolean;
+    mobileMenuOpen?: boolean;
+    onToggleMobileMenu?: (() => void) | null;
+    showMobileRightToggle?: boolean;
+    mobileRightOpen?: boolean;
+    onToggleMobileRight?: (() => void) | null;
+    downloadProgress?: {
+      count: number;
+      percentage: number;
+    } | null;
+  }
+
+  let {
+    showHome = true,
+    onHome = null,
+    showSidebarToggle = false,
+    sidebarVisible = true,
+    onToggleSidebar = null,
+    showMobileMenuToggle = false,
+    mobileMenuOpen = false,
+    onToggleMobileMenu = null,
+    showMobileRightToggle = false,
+    mobileRightOpen = false,
+    onToggleMobileRight = null,
+    downloadProgress = null,
+  }: Props = $props();
 
   function handleHome(): void {
     if (onHome) {
@@ -259,5 +276,26 @@
       {/if}
       <span class="hidden sm:inline">Downloads</span>
     </a>
+    <a
+      href="/#/integrations"
+      class="text-xs md:text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-1.5 md:gap-2 cursor-pointer"
+      title="Integration configs for external tools"
+    >
+      <svg
+        class="w-4 h-4"
+        viewBox="0 0 24 24"
+        fill="none"
+        stroke="currentColor"
+        stroke-width="2"
+        stroke-linecap="round"
+        stroke-linejoin="round"
+      >
+        <path d="M10 13a5 5 0 0 0 7.54.54l3-3a5 5 0 0 0-7.07-7.07l-1.72 1.71" />
+        <path
+          d="M14 11a5 5 0 0 0-7.54-.54l-3 3a5 5 0 0 0 7.07 7.07l1.71-1.71"
+        />
+      </svg>
+      <span class="hidden sm:inline">Integrations</span>
+    </a>
   </nav>
 </header>

dashboard/src/lib/components/IntegrationCard.svelte

@@ -0,0 +1,62 @@
+<script lang="ts">
+  import { copyText } from "$lib/utils/clipboard";
+
+  interface Props {
+    title: string;
+    subtitle: string;
+    config: string;
+    description?: string;
+    language?: "json" | "bash";
+  }
+
+  let {
+    title,
+    subtitle,
+    config,
+    description = "",
+    language = "json",
+  }: Props = $props();
+
+  let copied = $state(false);
+  let failed = $state(false);
+
+  async function copyToClipboard() {
+    const ok = await copyText(config);
+    if (ok) {
+      copied = true;
+      setTimeout(() => (copied = false), 2000);
+    } else {
+      failed = true;
+      setTimeout(() => (failed = false), 2000);
+    }
+  }
+</script>
+
+<div
+  class="border border-exo-light-gray/20 rounded-lg bg-exo-medium-gray/20 overflow-hidden"
+>
+  <div class="flex items-center justify-between px-5 py-4">
+    <div>
+      <h3 class="text-white text-sm font-semibold tracking-wide">{title}</h3>
+      <p class="text-exo-light-gray/60 text-xs mt-0.5 font-mono">{subtitle}</p>
+    </div>
+    <button
+      onclick={copyToClipboard}
+      class="px-3 py-1.5 text-xs rounded border transition-all duration-200 cursor-pointer
+        {copied
+        ? 'border-green-500/50 text-green-400 bg-green-500/10'
+        : failed
+          ? 'border-red-500/50 text-red-400 bg-red-500/10'
+          : 'border-exo-light-gray/30 text-exo-light-gray hover:border-exo-yellow/50 hover:text-exo-yellow'}"
+    >
+      {copied ? "Copied!" : failed ? "Copy failed" : "Copy"}
+    </button>
+  </div>
+  {#if description}
+    <p class="text-exo-light-gray/70 text-xs px-5 pb-3">{description}</p>
+  {/if}
+  <div class="bg-black/30 border-t border-exo-light-gray/10">
+    <pre
+      class="text-xs text-exo-light-gray/90 font-mono p-4 overflow-x-auto whitespace-pre">{config}</pre>
+  </div>
+</div>

dashboard/src/lib/components/ModelCard.svelte

@@ -16,7 +16,9 @@
       perNode?: Array<{
         nodeId: string;
         nodeName: string;
-        progress: DownloadProgress;
+        status: "completed" | "partial" | "pending" | "downloading";
+        percentage: number;
+        progress: DownloadProgress | null;
       }>;
     } | null;
     nodes?: Record<string, NodeInfo>;
@@ -145,10 +147,7 @@
     return `${s}s`;
   }
 
-  const isDownloading = $derived(downloadStatus?.isDownloading ?? false);
-  const progress = $derived(downloadStatus?.progress);
-  const percentage = $derived(progress?.percentage ?? 0);
-  let expandedNodes = $state<Set<string>>(new Set());
+  const perNode = $derived(downloadStatus?.perNode ?? []);
 
   function toggleNodeDetails(nodeId: string): void {
     const next = new Set(expandedNodes);
@@ -587,23 +586,49 @@
       </span>
     </div>
 
-    <!-- Download Status -->
-    {#if isDownloading && progress}
+    <!-- Download Status (per-node) -->
+    {#if perNode.length > 0}
       <div class="mb-2 space-y-1">
-        <div class="flex items-center justify-between text-xs font-mono">
-          <span class="text-blue-400 tracking-wider uppercase">Downloading</span
-          >
-          <span class="text-white/60"
-            >{percentage.toFixed(1)}% &middot; {formatSpeed(progress.speed)}
-            &middot; {formatEta(progress.etaMs)}</span
-          >
-        </div>
-        <div class="h-1 bg-exo-medium-gray/30 rounded overflow-hidden">
-          <div
-            class="h-full bg-blue-500/70 transition-all duration-300"
-            style="width: {percentage}%"
-          ></div>
+        <div
+          class="text-[10px] font-mono text-white/20 tracking-widest uppercase"
+        >
+          Download progress
         </div>
+        {#each perNode as node}
+          <div class="flex items-center gap-2 text-xs font-mono">
+            <span class="text-white/40 w-20 truncate" title={node.nodeId}
+              >{node.nodeName}</span
+            >
+            <div
+              class="flex-1 h-1 bg-exo-medium-gray/30 rounded overflow-hidden"
+            >
+              <div
+                class="h-full transition-all duration-300 {node.status ===
+                'downloading'
+                  ? 'bg-blue-500/70'
+                  : node.status === 'completed'
+                    ? 'bg-exo-yellow/40'
+                    : 'bg-white/20'}"
+                style="width: {node.percentage}%"
+              ></div>
+            </div>
+            <span
+              class="text-right {node.status === 'completed'
+                ? 'text-exo-yellow/60'
+                : node.status === 'downloading'
+                  ? 'text-blue-400/60'
+                  : 'text-white/30'}"
+            >
+              {#if node.status === "downloading" && node.progress}
+                {Math.round(node.percentage)}% {formatSpeed(
+                  node.progress.speed,
+                )}
+              {:else}
+                {node.percentage > 0 ? `${Math.round(node.percentage)}%` : "0%"}
+              {/if}
+            </span>
+          </div>
+        {/each}
       </div>
     {/if}
 
@@ -662,15 +687,7 @@
             {@const allConnections =
               isDebugMode && usedNodes.length > 1
                 ? (() => {
-                    const conns: Array<{
-                      ip: string;
-                      iface: string | null;
-                      from: string;
-                      to: string;
-                      midX: number;
-                      midY: number;
-                      arrow: string;
-                    }> = [];
+                    const conns: Array = [];
                     for (let i = 0; i < usedNodes.length; i++) {
                       for (let j = i + 1; j < usedNodes.length; j++) {
                         const n1 = usedNodes[i];
@@ -682,7 +699,12 @@
                           const toPos = nodePositions[c.to];
                           const arrow =
                             fromPos && toPos ? getArrow(fromPos, toPos) : "→";
-                          conns.push({ ...c, midX, midY, arrow });
+                          conns.push({
+                            ...c,
+                            midX,
+                            midY,
+                            arrow,
+                          });
                         }
                       }
                     }

dashboard/src/lib/components/ModelFilterPopover.svelte

@@ -73,7 +73,7 @@
 
 <!-- svelte-ignore a11y_no_static_element_interactions -->
 <div
-  class="filter-popover absolute right-0 top-full mt-2 w-64 bg-exo-dark-gray border border-exo-yellow/10 rounded-lg shadow-xl z-10"
+  class="filter-popover absolute right-0 top-full mt-2 w-64 bg-exo-dark-gray border border-exo-yellow/10 rounded-lg shadow-xl z-20"
   transition:fly={{ y: -10, duration: 200, easing: cubicOut }}
   onclick={(e) => e.stopPropagation()}
   role="dialog"

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -457,8 +457,10 @@
       "deepseek",
       "gpt-oss",
       "llama",
+      "gemma",
       "flux",
       "qwen-image",
+      "nemotron",
     ];
     return Array.from(families).sort((a, b) => {
       const aIdx = familyOrder.indexOf(a);

dashboard/src/lib/components/index.ts

@@ -13,3 +13,4 @@ export { default as ModelFilterPopover } from "./ModelFilterPopover.svelte";
 export { default as ModelPickerGroup } from "./ModelPickerGroup.svelte";
 export { default as ModelPickerModal } from "./ModelPickerModal.svelte";
 export { default as ChatModelSelector } from "./ChatModelSelector.svelte";
+export { default as IntegrationCard } from "./IntegrationCard.svelte";

dashboard/src/lib/stores/app.svelte.ts

@@ -257,11 +257,12 @@ interface RawStateResponse {
 }
 
 export interface MessageAttachment {
-  type: "image" | "text" | "file" | "generated-image";
+  type: "image" | "text" | "file" | "generated-image" | "pdf";
   name: string;
   content?: string;
   preview?: string;
   mimeType?: string;
+  pageImages?: string[];
 }
 
 export interface TopLogprob {
@@ -1793,6 +1794,14 @@ class AppStore {
             this.persistConversation(targetConversationId);
           }
         },
+        {
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+            }
+          },
+        },
       );
 
       // Final update
@@ -1990,6 +1999,14 @@ class AppStore {
             this.persistConversation(targetConversationId);
           }
         },
+        {
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+            }
+          },
+        },
       );
 
       // Final cleanup of the message (if conversation still exists)
@@ -2226,6 +2243,7 @@ class AppStore {
       type: string;
       textContent?: string;
       preview?: string;
+      pageImages?: string[];
     }[],
     enableThinking?: boolean | null,
   ): Promise<void> {
@@ -2261,6 +2279,20 @@ class AppStore {
             preview: file.preview,
             mimeType: file.type,
           });
+        } else if (
+          file.pageImages ||
+          (file.textContent && file.type === "application/pdf")
+        ) {
+          attachments.push({
+            type: "pdf",
+            name: file.name,
+            content: file.textContent,
+            pageImages: file.pageImages,
+            mimeType: file.type,
+          });
+          if (file.textContent) {
+            fileContext += `\n\n[File: ${file.name}]\n\`\`\`\n${file.textContent}\n\`\`\``;
+          }
         } else if (file.textContent) {
           attachments.push({
             type: "text",
@@ -2328,13 +2360,70 @@ class AppStore {
       const apiMessages = [
         systemPrompt,
         ...targetConversation.messages.slice(0, -1).map((m) => {
-          // Build content including any text file attachments
+          // Check if this message has image or PDF attachments
+          const visualAttachments = m.attachments?.filter(
+            (a) =>
+              (a.type === "image" && a.preview) ||
+              (a.type === "pdf" && a.pageImages?.length),
+          );
+
+          if (visualAttachments && visualAttachments.length > 0) {
+            // Build multimodal content array (OpenAI vision format)
+            const contentParts: Array<
+              | { type: "text"; text: string }
+              | { type: "image_url"; image_url: { url: string } }
+            > = [];
+
+            // Add image parts first
+            for (const att of visualAttachments) {
+              if (att.type === "image" && att.preview) {
+                contentParts.push({
+                  type: "image_url",
+                  image_url: { url: att.preview },
+                });
+              } else if (att.type === "pdf" && att.pageImages) {
+                for (const pageImg of att.pageImages) {
+                  contentParts.push({
+                    type: "image_url",
+                    image_url: { url: pageImg },
+                  });
+                }
+              }
+            }
+
+            // Build text content including any text/pdf file attachments
+            let textContent = m.content;
+            if (m.attachments) {
+              for (const attachment of m.attachments) {
+                if (
+                  (attachment.type === "text" || attachment.type === "pdf") &&
+                  attachment.content
+                ) {
+                  textContent += `\n\n[File: ${attachment.name}]\n\`\`\`\n${attachment.content}\n\`\`\``;
+                }
+              }
+            }
+
+            if (textContent) {
+              contentParts.push({ type: "text", text: textContent });
+            }
+
+            return {
+              role: m.role,
+              content: contentParts,
+            };
+          }
+
+          // Text-only message (original path)
           let msgContent = m.content;
 
-          // Add text attachments as context
+          // Add text/pdf attachments as context
           if (m.attachments) {
             for (const attachment of m.attachments) {
-              if (attachment.type === "text" && attachment.content) {
+              if (
+                (attachment.type === "text" || attachment.type === "pdf") &&
+                attachment.content
+              ) {
                 msgContent += `\n\n[File: ${attachment.name}]\n\`\`\`\n${attachment.content}\n\`\`\``;
               }
             }
@@ -2397,7 +2486,7 @@ class AppStore {
 
       let streamedContent = "";
       let streamedThinking = "";
-
+      let serverTpsReceived = false;
       interface ChatCompletionChunk {
         choices?: Array<{
           delta?: { content?: string; reasoning_content?: string };
@@ -2462,7 +2551,6 @@ class AppStore {
             tokenCount += 1;
             this.totalTokens = tokenCount;
 
-            // Update real-time TPS during streaming
             if (firstTokenTime !== null && tokenCount > 1) {
               const elapsed = performance.now() - firstTokenTime;
               this.tps = (tokenCount / elapsed) * 1000;
@@ -2513,16 +2601,24 @@ class AppStore {
               startedAt: this.prefillProgress?.startedAt ?? performance.now(),
             };
           },
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+              serverTpsReceived = true;
+            }
+          },
         },
       );
 
       // Clear prefill progress after stream ends
       this.prefillProgress = null;
 
-      // Calculate final TPS
-      if (firstTokenTime !== null && tokenCount > 1) {
+      // Use server-side TPS if available, otherwise fall back to client-side
+      if (!serverTpsReceived && firstTokenTime !== null && tokenCount > 1) {
         const totalGenerationTime = performance.now() - firstTokenTime;
-        this.tps = (tokenCount / totalGenerationTime) * 1000; // tokens per second
+        this.tps = (tokenCount / totalGenerationTime) * 1000;
       }
 
       // Final cleanup of the message (if conversation still exists)
@@ -2627,6 +2723,9 @@ class AppStore {
     this.syncActiveMessagesIfNeeded(targetConversationId);
     this.saveConversationsToStorage();
 
+    const abortController = new AbortController();
+    this.currentAbortController = abortController;
+
     try {
       // Determine the model to use
       const model = this.getModelForRequest(modelId);
@@ -2681,6 +2780,7 @@ class AppStore {
           "Content-Type": "application/json",
         },
         body: JSON.stringify(requestBody),
+        signal: abortController.signal,
       });
 
       if (!response.ok) {
@@ -2820,14 +2920,27 @@ class AppStore {
         );
       }
     } catch (error) {
-      console.error("Error generating image:", error);
-      this.handleStreamingError(
-        error,
-        targetConversationId,
-        assistantMessage.id,
-        "Failed to generate image",
-      );
+      if (abortController.signal.aborted) {
+        this.updateConversationMessage(
+          targetConversationId,
+          assistantMessage.id,
+          (msg) => {
+            msg.content = "Cancelled";
+            msg.attachments = [];
+          },
+        );
+        this.syncActiveMessagesIfNeeded(targetConversationId);
+      } else {
+        console.error("Error generating image:", error);
+        this.handleStreamingError(
+          error,
+          targetConversationId,
+          assistantMessage.id,
+          "Failed to generate image",
+        );
+      }
     } finally {
+      this.currentAbortController = null;
       this.isLoading = false;
       this.saveConversationsToStorage();
     }
@@ -2891,6 +3004,9 @@ class AppStore {
     // Clear editing state
     this.editingImage = null;
 
+    const abortController = new AbortController();
+    this.currentAbortController = abortController;
+
     try {
       // Determine the model to use
       const model = this.getModelForRequest(modelId);
@@ -2952,6 +3068,7 @@ class AppStore {
       const apiResponse = await fetch("/v1/images/edits", {
         method: "POST",
         body: formData,
+        signal: abortController.signal,
       });
 
       if (!apiResponse.ok) {
@@ -3052,14 +3169,27 @@ class AppStore {
         );
       }
     } catch (error) {
-      console.error("Error editing image:", error);
-      this.handleStreamingError(
-        error,
-        targetConversationId,
-        assistantMessage.id,
-        "Failed to edit image",
-      );
+      if (abortController.signal.aborted) {
+        this.updateConversationMessage(
+          targetConversationId,
+          assistantMessage.id,
+          (msg) => {
+            msg.content = "Cancelled";
+            msg.attachments = [];
+          },
+        );
+        this.syncActiveMessagesIfNeeded(targetConversationId);
+      } else {
+        console.error("Error editing image:", error);
+        this.handleStreamingError(
+          error,
+          targetConversationId,
+          assistantMessage.id,
+          "Failed to edit image",
+        );
+      }
     } finally {
+      this.currentAbortController = null;
       this.isLoading = false;
       this.saveConversationsToStorage();
     }
@@ -3126,6 +3256,31 @@ class AppStore {
     }
   }
 
+  /**
+   * Cancel/pause an active download on a specific node
+   */
+  async cancelDownload(nodeId: string, modelId: string): Promise<void> {
+    try {
+      const response = await fetch("/download/cancel", {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({
+          targetNodeId: nodeId,
+          modelId: modelId,
+        }),
+      });
+      if (!response.ok) {
+        const errorText = await response.text();
+        throw new Error(
+          `Failed to cancel download: ${response.status} - ${errorText}`,
+        );
+      }
+    } catch (error) {
+      console.error("Error cancelling download:", error);
+      throw error;
+    }
+  }
+
   /**
    * Delete a downloaded model from a specific node
    */
@@ -3248,6 +3403,7 @@ export const sendMessage = (
     type: string;
     textContent?: string;
     preview?: string;
+    pageImages?: string[];
   }[],
   enableThinking?: boolean | null,
 ) => appStore.sendMessage(content, files, enableThinking);
@@ -3346,6 +3502,8 @@ export const resetImageGenerationParams = () =>
 // Download actions
 export const startDownload = (nodeId: string, shardMetadata: object) =>
   appStore.startDownload(nodeId, shardMetadata);
+export const cancelDownload = (nodeId: string, modelId: string) =>
+  appStore.cancelDownload(nodeId, modelId);
 export const deleteDownload = (nodeId: string, modelId: string) =>
   appStore.deleteDownload(nodeId, modelId);
 

dashboard/src/lib/types/files.ts

@@ -2,6 +2,41 @@
  * File attachment types for the chat interface
  */
 
+import { getDocument, GlobalWorkerOptions, version } from "pdfjs-dist";
+import type { DocumentInitParameters } from "pdfjs-dist/types/src/display/api";
+
+// Safari (through at least 18/26) does not implement
+// ReadableStream.prototype[Symbol.asyncIterator], which pdfjs-dist uses
+// internally in getTextContent(). Without this polyfill, `for await (const n
+// of stream)` throws "undefined is not a function" and PDF processing fails.
+if (
+  typeof ReadableStream !== "undefined" &&
+  !(ReadableStream.prototype as unknown as Record<symbol, unknown>)[
+    Symbol.asyncIterator
+  ]
+) {
+  (ReadableStream.prototype as unknown as Record<symbol, unknown>)[
+    Symbol.asyncIterator
+  ] = async function* (this: ReadableStream<unknown>) {
+    const reader = this.getReader();
+    try {
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) return;
+        yield value;
+      }
+    } finally {
+      reader.releaseLock();
+    }
+  };
+}
+
+GlobalWorkerOptions.workerSrc = `https://cdn.jsdelivr.net/npm/pdfjs-dist@${version}/build/pdf.worker.mjs`;
+
+const PDF_PAGE_SCALE = 2.0;
+const PDF_MAX_PAGES = 20;
+const PDF_MAX_TEXT_CHARS = 100_000;
+
 export interface ChatUploadedFile {
   id: string;
   name: string;
@@ -10,6 +45,7 @@ export interface ChatUploadedFile {
   file: File;
   preview?: string;
   textContent?: string;
+  pageImages?: string[];
 }
 
 export interface ChatAttachment {
@@ -194,6 +230,58 @@ export function readFileAsText(file: File): Promise<string> {
   });
 }
 
+async function extractPdfContent(
+  file: File,
+): Promise<{ text: string; pageImages: string[] }> {
+  const arrayBuffer = await file.arrayBuffer();
+  const pdf = await getDocument({
+    data: new Uint8Array(arrayBuffer),
+    useSystemFonts: true,
+  } as DocumentInitParameters).promise;
+
+  const numPages = Math.min(pdf.numPages, PDF_MAX_PAGES);
+  const pageTexts: string[] = [];
+  const pageImages: string[] = [];
+
+  for (let i = 1; i <= numPages; i++) {
+    const page = await pdf.getPage(i);
+
+    const content = await page.getTextContent();
+    const strings = content.items
+      .filter((item: any) => "str" in item)
+      .map((item: any) => item.str as string);
+    pageTexts.push(strings.join(" "));
+
+    const viewport = page.getViewport({ scale: PDF_PAGE_SCALE });
+    const canvas = new OffscreenCanvas(viewport.width, viewport.height);
+    const ctx = canvas.getContext("2d");
+    if (ctx) {
+      await page.render({ canvasContext: ctx as any, viewport }).promise;
+      const blob = await canvas.convertToBlob({
+        type: "image/jpeg",
+        quality: 0.8,
+      });
+      const reader = new FileReader();
+      const dataUrl = await new Promise<string>((resolve, reject) => {
+        reader.onload = () => resolve(reader.result as string);
+        reader.onerror = () => reject(reader.error);
+        reader.readAsDataURL(blob);
+      });
+      pageImages.push(dataUrl);
+    }
+  }
+
+  let text = pageTexts.join("\n\n").trim();
+  if (text.length > PDF_MAX_TEXT_CHARS) {
+    text = text.slice(0, PDF_MAX_TEXT_CHARS) + "\n\n[truncated]";
+  }
+  if (pdf.numPages > PDF_MAX_PAGES) {
+    text += `\n\n[showing ${PDF_MAX_PAGES} of ${pdf.numPages} pages]`;
+  }
+
+  return { text, pageImages };
+}
+
 /**
  * Process uploaded files into ChatUploadedFile format
  */
@@ -223,7 +311,12 @@ export async function processUploadedFiles(
         const textContent = await readFileAsText(file);
         results.push({ ...base, textContent });
       } else if (category === "pdf") {
-        results.push(base);
+        const { text, pageImages } = await extractPdfContent(file);
+        results.push({
+          ...base,
+          textContent: text || undefined,
+          pageImages: pageImages.length > 0 ? pageImages : undefined,
+        });
       } else if (category === "audio") {
         const preview = await readFileAsDataURL(file);
         results.push({ ...base, preview });

dashboard/src/lib/utils/clipboard.ts

@@ -0,0 +1,55 @@
+export async function copyText(text: string): Promise<boolean> {
+  if (
+    typeof window !== "undefined" &&
+    window.isSecureContext &&
+    navigator.clipboard?.writeText
+  ) {
+    try {
+      await navigator.clipboard.writeText(text);
+      return true;
+    } catch {
+      // fall through to execCommand fallback
+    }
+  }
+
+  if (typeof document === "undefined") {
+    return false;
+  }
+
+  const textarea = document.createElement("textarea");
+  textarea.value = text;
+  textarea.setAttribute("readonly", "");
+  textarea.style.position = "fixed";
+  textarea.style.top = "0";
+  textarea.style.left = "0";
+  textarea.style.width = "1px";
+  textarea.style.height = "1px";
+  textarea.style.padding = "0";
+  textarea.style.border = "none";
+  textarea.style.outline = "none";
+  textarea.style.boxShadow = "none";
+  textarea.style.background = "transparent";
+  textarea.style.opacity = "0";
+  document.body.appendChild(textarea);
+
+  const previousSelection = document.getSelection();
+  const previousRange =
+    previousSelection && previousSelection.rangeCount > 0
+      ? previousSelection.getRangeAt(0)
+      : null;
+
+  try {
+    textarea.focus();
+    textarea.select();
+    textarea.setSelectionRange(0, text.length);
+    return document.execCommand("copy");
+  } catch {
+    return false;
+  } finally {
+    document.body.removeChild(textarea);
+    if (previousRange && previousSelection) {
+      previousSelection.removeAllRanges();
+      previousSelection.addRange(previousRange);
+    }
+  }
+}

dashboard/src/routes/+page.svelte

@@ -42,6 +42,7 @@
     setSelectedChatModel,
     selectedChatModel,
     sendMessage,
+    thinkingEnabled,
     generateImage,
     editImage,
     editingImage,
@@ -264,6 +265,7 @@
 
   let mounted = $state(false);
   let localNodeId = $state<string | null>(null);
+  let pendingFirefoxQuery = $state<string | null>(null); // ?q= param deferred until state loads
 
   // ── Onboarding wizard state ──
   const ONBOARDING_COMPLETE_KEY = "exo-onboarding-complete";
@@ -852,7 +854,7 @@
   ) {
     const model = selectedChatModel();
     if (!model) {
-      sendMessage(content, files, null);
+      sendMessage(content, files, thinkingEnabled());
       return;
     }
 
@@ -880,7 +882,7 @@
     }
 
     // Default: text chat
-    sendMessage(content, files, null);
+    sendMessage(content, files, thinkingEnabled());
   }
 
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
@@ -1308,6 +1310,20 @@
       return;
     }
 
+    // Firefox AI sidebar integration: handle ?q= query parameter
+    // Firefox's built-in AI sidebar (about:config: browser.ml.chat.enabled) sends
+    // the user's prompt as ?q=<URL-encoded prompt> to the configured provider URL.
+    // See: https://support.mozilla.org/en-US/kb/ai-chatbot
+    const queryParam = params.get("q");
+    if (queryParam) {
+      // Clean up the URL to prevent re-submission on page refresh
+      window.history.replaceState({}, "", window.location.pathname);
+      // Defer the auto-send until cluster state is loaded (topologyData,
+      // instances, availableMemory) so that model auto-selection works
+      // correctly. The $effect below will pick this up once data is ready.
+      pendingFirefoxQuery = queryParam;
+    }
+
     // Check server-side onboarding state (persisted in ~/.exo)
     try {
       const res = await fetch("/onboarding");
@@ -1328,6 +1344,18 @@
     }
   });
 
+  // Deferred Firefox AI sidebar auto-send: wait for cluster state and model
+  // list before submitting. Both data (from /state polling) and models (from
+  // the async /models fetch in onMount) must be loaded for handleAutoSend to
+  // correctly auto-select a model.
+  $effect(() => {
+    if (pendingFirefoxQuery && data && models.length > 0) {
+      const query = pendingFirefoxQuery;
+      pendingFirefoxQuery = null;
+      handleChatSend(query);
+    }
+  });
+
   async function fetchModels() {
     try {
       const response = await fetch("/models");
@@ -1535,34 +1563,44 @@
   }
 
   // Helper to get download status for a model (checks all downloads for matching model ID)
-  function getModelDownloadStatus(modelId: string): {
+  type NodeDownloadStatus = {
+    nodeId: string;
+    nodeName: string;
+    status: "completed" | "partial" | "pending" | "downloading";
+    percentage: number;
+    progress: DownloadProgress | null;
+  };
+
+  // Shared helper: collect per-node download status for a model across a set of nodes.
+  // Handles deduplication, entry parsing, and aggregation in one place.
+  function collectDownloadStatus(
+    modelId: string,
+    nodeIds?: string[],
+  ): {
     isDownloading: boolean;
     progress: DownloadProgress | null;
-    perNode: Array<{
-      nodeId: string;
-      nodeName: string;
-      progress: DownloadProgress;
-    }>;
+    perNode: NodeDownloadStatus[];
+    failedError: string | null;
   } {
+    const empty = {
+      isDownloading: false,
+      progress: null,
+      perNode: [] as NodeDownloadStatus[],
+      failedError: null,
+    };
+
     if (!downloadsData || Object.keys(downloadsData).length === 0) {
-      return { isDownloading: false, progress: null, perNode: [] };
+      return empty;
     }
 
-    let totalBytes = 0;
-    let downloadedBytes = 0;
-    let totalSpeed = 0;
-    let completedFiles = 0;
-    let totalFiles = 0;
-    let isDownloading = false;
-    const allFiles: DownloadProgress["files"] = [];
-    const perNode: Array<{
-      nodeId: string;
-      nodeName: string;
-      progress: DownloadProgress;
-    }> = [];
+    // Deduplicate by nodeId — a node can have multiple entries for the same model
+    // (e.g. PipelineShardMetadata + TensorShardMetadata). Keep the last entry,
+    // which is the most recently applied event.
+    const perNodeMap = new Map<string, NodeDownloadStatus>();
 
-    // Check all nodes for downloads matching this model
+    const nodeIdSet = nodeIds ? new Set(nodeIds) : null;
     for (const [nodeId, nodeDownloads] of Object.entries(downloadsData)) {
+      if (nodeIdSet && !nodeIdSet.has(nodeId)) continue;
       if (!Array.isArray(nodeDownloads)) continue;
 
       for (const downloadWrapped of nodeDownloads) {
@@ -1575,29 +1613,45 @@
         const downloadPayload = (downloadWrapped as Record<string, unknown>)[
           downloadKind
         ] as Record<string, unknown>;
-
-        if (
-          downloadKind !== "DownloadOngoing" &&
-          downloadKind !== "DownloadPending"
-        )
-          continue;
         if (!downloadPayload) continue;
 
         const downloadModelId = extractModelIdFromDownload(downloadPayload);
+        if (!downloadModelId || downloadModelId !== modelId) continue;
 
-        // Match if the model ID contains or equals the requested model
-        // (handles cases like "mlx-community/Meta-Llama..." matching)
-        if (
-          !downloadModelId ||
-          !downloadModelId.includes(modelId.split("/").pop() || modelId)
-        ) {
-          // Try exact match or partial match
-          if (downloadModelId !== modelId) continue;
+        // DownloadFailed — return with any data collected so far
+        if (downloadKind === "DownloadFailed") {
+          return {
+            isDownloading: false,
+            progress: null,
+            perNode: Array.from(perNodeMap.values()),
+            failedError:
+              (downloadPayload.errorMessage as string) ||
+              (downloadPayload.error_message as string) ||
+              "Download failed",
+          };
+        }
+
+        if (
+          downloadKind !== "DownloadOngoing" &&
+          downloadKind !== "DownloadPending" &&
+          downloadKind !== "DownloadCompleted"
+        )
+          continue;
+
+        const nodeName =
+          data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
+
+        if (downloadKind === "DownloadCompleted") {
+          perNodeMap.set(nodeId, {
+            nodeId,
+            nodeName,
+            status: "completed",
+            percentage: 100,
+            progress: null,
+          });
+          continue;
         }
 
-        // For DownloadPending with partial bytes (paused/resumed downloads),
-        // synthesize a progress object from the top-level downloaded/total fields
-        let progress: DownloadProgress | null;
         if (downloadKind === "DownloadPending") {
           const pendingDownloaded = getBytes(
             downloadPayload.downloaded ??
@@ -1610,44 +1664,67 @@
               downloadPayload.totalBytes,
           );
           if (pendingDownloaded <= 0 && pendingTotal <= 0) continue;
-          isDownloading = true;
-          progress = {
-            totalBytes: pendingTotal,
-            downloadedBytes: pendingDownloaded,
-            speed: 0,
-            etaMs: 0,
-            percentage:
-              pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0,
-            completedFiles: 0,
-            totalFiles: 0,
-            files: [],
-          };
-        } else {
-          isDownloading = true;
-          progress = parseDownloadProgress(downloadPayload);
+          const pct =
+            pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0;
+          perNodeMap.set(nodeId, {
+            nodeId,
+            nodeName,
+            status: pendingDownloaded > 0 ? "partial" : "pending",
+            percentage: pct,
+            progress: null,
+          });
+          continue;
         }
 
-        if (progress) {
-          // Sum all values across nodes - each node downloads independently
-          totalBytes += progress.totalBytes;
-          downloadedBytes += progress.downloadedBytes;
-          totalSpeed += progress.speed;
-          completedFiles += progress.completedFiles;
-          totalFiles += progress.totalFiles;
-          allFiles.push(...progress.files);
+        // DownloadOngoing
+        const progress = parseDownloadProgress(downloadPayload);
+        if (
+          !progress ||
+          (progress.downloadedBytes <= 0 && progress.totalBytes <= 0)
+        )
+          continue;
 
-          const nodeName =
-            data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
-          perNode.push({ nodeId, nodeName, progress });
-        }
+        perNodeMap.set(nodeId, {
+          nodeId,
+          nodeName,
+          status: "downloading",
+          percentage: progress.percentage,
+          progress,
+        });
+      }
+    }
+
+    // Aggregate from deduplicated per-node entries
+    const perNode = Array.from(perNodeMap.values());
+    let totalBytes = 0;
+    let downloadedBytes = 0;
+    let totalSpeed = 0;
+    let completedFiles = 0;
+    let totalFiles = 0;
+    let isDownloading = false;
+    const allFiles: DownloadProgress["files"] = [];
+
+    for (const node of perNode) {
+      if (node.status === "downloading" && node.progress) {
+        isDownloading = true;
+        totalBytes += node.progress.totalBytes;
+        downloadedBytes += node.progress.downloadedBytes;
+        totalSpeed += node.progress.speed;
+        completedFiles += node.progress.completedFiles;
+        totalFiles += node.progress.totalFiles;
+        allFiles.push(...node.progress.files);
       }
     }
 
     if (!isDownloading) {
-      return { isDownloading: false, progress: null, perNode: [] };
+      return {
+        isDownloading: false,
+        progress: null,
+        perNode,
+        failedError: null,
+      };
     }
 
-    // ETA = total remaining bytes / total speed across all nodes
     const remainingBytes = totalBytes - downloadedBytes;
     const etaMs = totalSpeed > 0 ? (remainingBytes / totalSpeed) * 1000 : 0;
 
@@ -1664,9 +1741,21 @@
         files: allFiles,
       },
       perNode,
+      failedError: null,
     };
   }
 
+  function getModelDownloadStatus(
+    modelId: string,
+    nodeIds?: string[],
+  ): {
+    isDownloading: boolean;
+    progress: DownloadProgress | null;
+    perNode: NodeDownloadStatus[];
+  } {
+    return collectDownloadStatus(modelId, nodeIds);
+  }
+
   // Helper to get download status for an instance
   function getInstanceDownloadStatus(
     instanceId: string,
@@ -1677,26 +1766,9 @@
     errorMessage: string | null;
     progress: DownloadProgress | null;
     statusText: string;
-    perNode: Array<{
-      nodeId: string;
-      nodeName: string;
-      progress: DownloadProgress;
-    }>;
+    perNode: NodeDownloadStatus[];
   } {
-    if (!downloadsData || Object.keys(downloadsData).length === 0) {
-      // No download data yet — defer to runner status instead of assuming RUNNING
-      const statusInfo = deriveInstanceStatus(instanceWrapped);
-      return {
-        isDownloading: false,
-        isFailed: false,
-        errorMessage: null,
-        progress: null,
-        statusText: statusInfo.statusText,
-        perNode: [],
-      };
-    }
-
-    // Unwrap the instance
+    // Unwrap the instance to get shard assignments
     const [instanceTag, instance] = getTagged(instanceWrapped);
     if (!instance || typeof instance !== "object") {
       return {
@@ -1716,132 +1788,9 @@
         modelId?: string;
       };
     };
-    const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
-    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
     const instanceModelId = inst.shardAssignments?.modelId;
 
-    // Build reverse mapping: runnerId -> nodeId
-    const runnerToNode: Record<string, string> = {};
-    for (const [nodeId, runnerId] of Object.entries(nodeToRunner)) {
-      runnerToNode[runnerId] = nodeId;
-    }
-
-    let totalBytes = 0;
-    let downloadedBytes = 0;
-    let totalSpeed = 0;
-    let completedFiles = 0;
-    let totalFiles = 0;
-    let isDownloading = false;
-    const allFiles: DownloadProgress["files"] = [];
-    const perNode: Array<{
-      nodeId: string;
-      nodeName: string;
-      progress: DownloadProgress;
-    }> = [];
-
-    // Check downloads for nodes that are part of this instance
-    for (const runnerId of Object.keys(runnerToShard)) {
-      const nodeId = runnerToNode[runnerId];
-      if (!nodeId) continue;
-
-      const nodeDownloads = downloadsData[nodeId];
-      if (!Array.isArray(nodeDownloads)) continue;
-
-      for (const downloadWrapped of nodeDownloads) {
-        if (!downloadWrapped || typeof downloadWrapped !== "object") continue;
-
-        const keys = Object.keys(downloadWrapped as Record<string, unknown>);
-        if (keys.length !== 1) continue;
-
-        const downloadKind = keys[0];
-        const downloadPayload = (downloadWrapped as Record<string, unknown>)[
-          downloadKind
-        ] as Record<string, unknown>;
-
-        // Handle DownloadFailed - return immediately with error info
-        if (downloadKind === "DownloadFailed") {
-          const downloadModelId = extractModelIdFromDownload(downloadPayload);
-          if (
-            instanceModelId &&
-            downloadModelId &&
-            downloadModelId === instanceModelId
-          ) {
-            return {
-              isDownloading: false,
-              isFailed: true,
-              errorMessage:
-                (downloadPayload.errorMessage as string) || "Download failed",
-              progress: null,
-              statusText: "FAILED",
-              perNode: [],
-            };
-          }
-        }
-
-        if (
-          downloadKind !== "DownloadOngoing" &&
-          downloadKind !== "DownloadPending"
-        )
-          continue;
-        if (!downloadPayload) continue;
-
-        // Check if this download is for this instance's model
-        const downloadModelId = extractModelIdFromDownload(downloadPayload);
-        if (
-          instanceModelId &&
-          downloadModelId &&
-          downloadModelId === instanceModelId
-        ) {
-          // For DownloadPending with partial bytes, synthesize progress
-          let progress: DownloadProgress | null;
-          if (downloadKind === "DownloadPending") {
-            const pendingDownloaded = getBytes(
-              downloadPayload.downloaded ??
-                downloadPayload.downloaded_bytes ??
-                downloadPayload.downloadedBytes,
-            );
-            const pendingTotal = getBytes(
-              downloadPayload.total ??
-                downloadPayload.total_bytes ??
-                downloadPayload.totalBytes,
-            );
-            if (pendingDownloaded <= 0 && pendingTotal <= 0) continue;
-            isDownloading = true;
-            progress = {
-              totalBytes: pendingTotal,
-              downloadedBytes: pendingDownloaded,
-              speed: 0,
-              etaMs: 0,
-              percentage:
-                pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0,
-              completedFiles: 0,
-              totalFiles: 0,
-              files: [],
-            };
-          } else {
-            isDownloading = true;
-            progress = parseDownloadProgress(downloadPayload);
-          }
-
-          if (progress) {
-            // Sum all values across nodes - each node downloads independently
-            totalBytes += progress.totalBytes;
-            downloadedBytes += progress.downloadedBytes;
-            totalSpeed += progress.speed;
-            completedFiles += progress.completedFiles;
-            totalFiles += progress.totalFiles;
-            allFiles.push(...progress.files);
-
-            const nodeName =
-              data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
-            perNode.push({ nodeId, nodeName, progress });
-          }
-        }
-      }
-    }
-
-    if (!isDownloading) {
-      // Check runner status for other states
+    if (!instanceModelId) {
       const statusInfo = deriveInstanceStatus(instanceWrapped);
       return {
         isDownloading: false,
@@ -1853,26 +1802,49 @@
       };
     }
 
-    // ETA = total remaining bytes / total speed across all nodes
-    const remainingBytes = totalBytes - downloadedBytes;
-    const etaMs = totalSpeed > 0 ? (remainingBytes / totalSpeed) * 1000 : 0;
+    // Get node IDs assigned to this instance
+    const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
+    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
+    const runnerToNode: Record<string, string> = {};
+    for (const [nodeId, runnerId] of Object.entries(nodeToRunner)) {
+      runnerToNode[runnerId] = nodeId;
+    }
+    const instanceNodeIds = Object.keys(runnerToShard)
+      .map((runnerId) => runnerToNode[runnerId])
+      .filter(Boolean);
+
+    const result = collectDownloadStatus(instanceModelId, instanceNodeIds);
+
+    if (result.failedError) {
+      return {
+        isDownloading: false,
+        isFailed: true,
+        errorMessage: result.failedError,
+        progress: null,
+        statusText: "FAILED",
+        perNode: [],
+      };
+    }
+
+    if (!result.isDownloading) {
+      const statusInfo = deriveInstanceStatus(instanceWrapped);
+      return {
+        isDownloading: false,
+        isFailed: statusInfo.statusText === "FAILED",
+        errorMessage: null,
+        progress: null,
+        statusText: statusInfo.statusText,
+        perNode: result.perNode,
+      };
+    }
 
     return {
       isDownloading: true,
       isFailed: false,
       errorMessage: null,
-      progress: {
-        totalBytes,
-        downloadedBytes,
-        speed: totalSpeed,
-        etaMs,
-        percentage: totalBytes > 0 ? (downloadedBytes / totalBytes) * 100 : 0,
-        completedFiles,
-        totalFiles,
-        files: allFiles,
-      },
+      progress: result.progress,
       statusText: "DOWNLOADING",
-      perNode,
+      perNode: result.perNode,
     };
   }
 
@@ -4630,7 +4602,7 @@
                 type="button"
                 onclick={() => {
                   completeOnboarding();
-                  sendMessage(chip);
+                  sendMessage(chip, undefined, thinkingEnabled());
                 }}
                 class="px-4 py-2 rounded-full border border-white/10 bg-white/5 text-sm text-white/60 hover:bg-white/10 hover:text-white/80 hover:border-white/20 transition-all duration-200 cursor-pointer"
               >
@@ -5369,10 +5341,10 @@
                             <div
                               class="mt-2 space-y-2 max-h-48 overflow-y-auto pr-1"
                             >
-                              {#each downloadInfo.perNode as nodeProg}
+                              {#each downloadInfo.perNode.filter((n) => n.status === "downloading" && n.progress) as nodeProg}
                                 {@const nodePercent = Math.min(
                                   100,
-                                  Math.max(0, nodeProg.progress.percentage),
+                                  Math.max(0, nodeProg.percentage),
                                 )}
                                 {@const isExpanded =
                                   instanceDownloadExpandedNodes.has(
@@ -5428,15 +5400,17 @@
                                     >
                                       <span
                                         >{formatBytes(
-                                          nodeProg.progress.downloadedBytes,
+                                          nodeProg.progress?.downloadedBytes ??
+                                            0,
                                         )} / {formatBytes(
-                                          nodeProg.progress.totalBytes,
+                                          nodeProg.progress?.totalBytes ?? 0,
                                         )}</span
                                       >
                                       <span
-                                        >{formatSpeed(nodeProg.progress.speed)} •
-                                        ETA {formatEta(
-                                          nodeProg.progress.etaMs,
+                                        >{formatSpeed(
+                                          nodeProg.progress?.speed ?? 0,
+                                        )} • ETA {formatEta(
+                                          nodeProg.progress?.etaMs ?? 0,
                                         )}</span
                                       >
                                     </div>
@@ -5444,14 +5418,14 @@
 
                                   {#if isExpanded}
                                     <div class="mt-2 space-y-1.5">
-                                      {#if nodeProg.progress.files.length === 0}
+                                      {#if nodeProg.progress?.files ?? [].length === 0}
                                         <div
                                           class="text-[11px] font-mono text-exo-light-gray/70"
                                         >
                                           No file details reported.
                                         </div>
                                       {:else}
-                                        {#each nodeProg.progress.files as f}
+                                        {#each nodeProg.progress?.files ?? [] as f}
                                           {@const filePercent = Math.min(
                                             100,
                                             Math.max(0, f.percentage ?? 0),
@@ -5927,12 +5901,15 @@
                 )}
                 {@const allPreviews = filteredPreviews()}
                 {#if selectedModel && allPreviews.length > 0}
-                  {@const downloadStatus = getModelDownloadStatus(
-                    selectedModel.id,
-                  )}
                   {@const tags = modelTags()[selectedModel.id] || []}
                   <div class="space-y-3">
                     {#each allPreviews as apiPreview, i}
+                      {@const downloadStatus = getModelDownloadStatus(
+                        selectedModel.id,
+                        apiPreview.memory_delta_by_node
+                          ? Object.keys(apiPreview.memory_delta_by_node)
+                          : undefined,
+                      )}
                       <div
                         role="group"
                         onmouseenter={() => {
@@ -6120,7 +6097,7 @@
                           onclick={() => {
                             chatLaunchState = "idle";
                             selectedChatCategory = null;
-                            sendMessage(prompt);
+                            sendMessage(prompt, undefined, thinkingEnabled());
                           }}
                           class="text-left px-3 py-2.5 text-xs text-exo-light-gray hover:text-white font-mono rounded-lg border border-exo-medium-gray/30 hover:border-exo-yellow/30 bg-exo-dark-gray/30 hover:bg-exo-dark-gray/60 transition-all duration-200 cursor-pointer"
                         >
@@ -6503,10 +6480,10 @@
                               <div
                                 class="mt-2 space-y-2 max-h-48 overflow-y-auto pr-1"
                               >
-                                {#each downloadInfo.perNode as nodeProg}
+                                {#each downloadInfo.perNode.filter((n) => n.status === "downloading" && n.progress) as nodeProg}
                                   {@const nodePercent = Math.min(
                                     100,
-                                    Math.max(0, nodeProg.progress.percentage),
+                                    Math.max(0, nodeProg.percentage),
                                   )}
                                   {@const isExpanded =
                                     instanceDownloadExpandedNodes.has(
@@ -6565,16 +6542,17 @@
                                       >
                                         <span
                                           >{formatBytes(
-                                            nodeProg.progress.downloadedBytes,
+                                            nodeProg.progress
+                                              ?.downloadedBytes ?? 0,
                                           )} / {formatBytes(
-                                            nodeProg.progress.totalBytes,
+                                            nodeProg.progress?.totalBytes ?? 0,
                                           )}</span
                                         >
                                         <span
                                           >{formatSpeed(
-                                            nodeProg.progress.speed,
+                                            nodeProg.progress?.speed ?? 0,
                                           )} • ETA {formatEta(
-                                            nodeProg.progress.etaMs,
+                                            nodeProg.progress?.etaMs ?? 0,
                                           )}</span
                                         >
                                       </div>
@@ -6582,14 +6560,14 @@
 
                                     {#if isExpanded}
                                       <div class="mt-2 space-y-1.5">
-                                        {#if nodeProg.progress.files.length === 0}
+                                        {#if nodeProg.progress?.files ?? [].length === 0}
                                           <div
                                             class="text-[11px] font-mono text-exo-light-gray/70"
                                           >
                                             No file details reported.
                                           </div>
                                         {:else}
-                                          {#each nodeProg.progress.files as f}
+                                          {#each nodeProg.progress?.files ?? [] as f}
                                             {@const filePercent = Math.min(
                                               100,
                                               Math.max(0, f.percentage ?? 0),

dashboard/src/routes/downloads/+page.svelte

@@ -9,6 +9,7 @@
     refreshState,
     lastUpdate as lastUpdateStore,
     startDownload,
+    cancelDownload,
     deleteDownload,
   } from "$lib/stores/app.svelte";
   import {
@@ -349,6 +350,59 @@
   });
 </script>
 
+{#snippet trashIcon()}
+  <svg
+    class="w-5 h-5"
+    viewBox="0 0 20 20"
+    fill="none"
+    stroke="currentColor"
+    stroke-width="2"
+  >
+    <path
+      d="M4 6h12M8 6V4h4v2m1 0v10a1 1 0 01-1 1H8a1 1 0 01-1-1V6h6"
+      stroke-linecap="round"
+      stroke-linejoin="round"
+    ></path>
+  </svg>
+{/snippet}
+
+{#snippet downloadIcon(size?: string)}
+  <svg
+    class={size ?? "w-5 h-5"}
+    viewBox="0 0 20 20"
+    fill="none"
+    stroke="currentColor"
+    stroke-width="2"
+  >
+    <path
+      d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
+      stroke-linecap="round"
+      stroke-linejoin="round"
+    ></path>
+  </svg>
+{/snippet}
+
+{#snippet pauseIcon()}
+  <svg class="w-5 h-5" viewBox="0 0 20 20" fill="currentColor">
+    <path
+      fill-rule="evenodd"
+      d="M6 4h2v12H6V4zm6 0h2v12h-2V4z"
+      clip-rule="evenodd"
+    ></path>
+  </svg>
+{/snippet}
+
+{#snippet deleteButton(nodeId: string, modelId: string)}
+  <button
+    type="button"
+    class="text-white/50 hover:text-red-400 transition-colors cursor-pointer"
+    onclick={() => deleteDownload(nodeId, modelId)}
+    title="Delete from this node"
+  >
+    {@render trashIcon()}
+  </button>
+{/snippet}
+
 <div class="min-h-screen bg-exo-dark-gray text-white">
   <HeaderNav showHome={true} />
   <div class="max-w-7xl mx-auto px-4 lg:px-8 py-6 space-y-6">
@@ -486,27 +540,7 @@
                         <span class="text-xs text-white/70"
                           >{formatBytes(cell.totalBytes)}</span
                         >
-                        <button
-                          type="button"
-                          class="text-white/50 hover:text-red-400 transition-colors mt-0.5 cursor-pointer"
-                          onclick={() =>
-                            deleteDownload(col.nodeId, row.modelId)}
-                          title="Delete from this node"
-                        >
-                          <svg
-                            class="w-5 h-5"
-                            viewBox="0 0 20 20"
-                            fill="none"
-                            stroke="currentColor"
-                            stroke-width="2"
-                          >
-                            <path
-                              d="M4 6h12M8 6V4h4v2m1 0v10a1 1 0 01-1 1H8a1 1 0 01-1-1V6h6"
-                              stroke-linecap="round"
-                              stroke-linejoin="round"
-                            ></path>
-                          </svg>
-                        </button>
+                        {@render deleteButton(col.nodeId, row.modelId)}
                       </div>
                     {:else if cell.kind === "downloading"}
                       <div
@@ -533,6 +567,18 @@
                         <span class="text-[10px] text-white/70"
                           >{formatSpeed(cell.speed)}</span
                         >
+                        <div class="flex gap-1 mt-0.5">
+                          <button
+                            type="button"
+                            class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
+                            onclick={() =>
+                              cancelDownload(col.nodeId, row.modelId)}
+                            title="Pause download"
+                          >
+                            {@render pauseIcon()}
+                          </button>
+                          {@render deleteButton(col.nodeId, row.modelId)}
+                        </div>
                       </div>
                     {:else if cell.kind === "pending"}
                       <div
@@ -558,32 +604,24 @@
                               ).toFixed(1)}%"
                             ></div>
                           </div>
-                          {#if row.shardMetadata}
-                            <button
-                              type="button"
-                              class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
-                              onclick={() =>
-                                startDownload(col.nodeId, row.shardMetadata!)}
-                              title="Resume download on this node"
-                            >
-                              <svg
-                                class="w-5 h-5"
-                                viewBox="0 0 20 20"
-                                fill="none"
-                                stroke="currentColor"
-                                stroke-width="2"
+                          <div class="flex gap-1">
+                            {#if row.shardMetadata}
+                              <button
+                                type="button"
+                                class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
+                                onclick={() =>
+                                  startDownload(col.nodeId, row.shardMetadata!)}
+                                title="Resume download on this node"
                               >
-                                <path
-                                  d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
-                                  stroke-linecap="round"
-                                  stroke-linejoin="round"
-                                ></path>
-                              </svg>
-                            </button>
-                          {:else}
-                            <span class="text-white/50 text-[10px]">paused</span
-                            >
-                          {/if}
+                                {@render downloadIcon()}
+                              </button>
+                            {:else}
+                              <span class="text-white/50 text-[10px]"
+                                >paused</span
+                              >
+                            {/if}
+                            {@render deleteButton(col.nodeId, row.modelId)}
+                          </div>
                         {:else if row.shardMetadata}
                           <button
                             type="button"
@@ -592,19 +630,7 @@
                               startDownload(col.nodeId, row.shardMetadata!)}
                             title="Start download on this node"
                           >
-                            <svg
-                              class="w-6 h-6"
-                              viewBox="0 0 20 20"
-                              fill="none"
-                              stroke="currentColor"
-                              stroke-width="2"
-                            >
-                              <path
-                                d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
-                                stroke-linecap="round"
-                                stroke-linejoin="round"
-                              ></path>
-                            </svg>
+                            {@render downloadIcon("w-6 h-6")}
                           </button>
                         {:else}
                           <span class="text-white/40 text-sm">...</span>
@@ -626,29 +652,20 @@
                             clip-rule="evenodd"
                           ></path>
                         </svg>
-                        {#if row.shardMetadata}
-                          <button
-                            type="button"
-                            class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
-                            onclick={() =>
-                              startDownload(col.nodeId, row.shardMetadata!)}
-                            title="Retry download on this node"
-                          >
-                            <svg
-                              class="w-5 h-5"
-                              viewBox="0 0 20 20"
-                              fill="none"
-                              stroke="currentColor"
-                              stroke-width="2"
+                        <div class="flex gap-1">
+                          {#if row.shardMetadata}
+                            <button
+                              type="button"
+                              class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
+                              onclick={() =>
+                                startDownload(col.nodeId, row.shardMetadata!)}
+                              title="Retry download on this node"
                             >
-                              <path
-                                d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
-                                stroke-linecap="round"
-                                stroke-linejoin="round"
-                              ></path>
-                            </svg>
-                          </button>
-                        {/if}
+                              {@render downloadIcon()}
+                            </button>
+                          {/if}
+                          {@render deleteButton(col.nodeId, row.modelId)}
+                        </div>
                       </div>
                     {:else}
                       <div
@@ -666,19 +683,7 @@
                               startDownload(col.nodeId, row.shardMetadata!)}
                             title="Download to this node"
                           >
-                            <svg
-                              class="w-5 h-5"
-                              viewBox="0 0 20 20"
-                              fill="none"
-                              stroke="currentColor"
-                              stroke-width="2"
-                            >
-                              <path
-                                d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
-                                stroke-linecap="round"
-                                stroke-linejoin="round"
-                              ></path>
-                            </svg>
+                            {@render downloadIcon()}
                           </button>
                         {/if}
                       </div>

dashboard/src/routes/integrations/+page.svelte

@@ -0,0 +1,577 @@
+<script lang="ts">
+  import { browser } from "$app/environment";
+  import { fade } from "svelte/transition";
+  import HeaderNav from "$lib/components/HeaderNav.svelte";
+  import IntegrationCard from "$lib/components/IntegrationCard.svelte";
+  import { instances, refreshState } from "$lib/stores/app.svelte";
+  import { onMount } from "svelte";
+
+  const apiUrl = browser
+    ? window.location.origin.replace("localhost", "127.0.0.1")
+    : "http://127.0.0.1:52415";
+
+  const instancesData = $derived(instances());
+
+  let modelCapabilities = $state<Record<string, string[]>>({});
+  let modelContextLengths = $state<Record<string, number>>({});
+
+  const runningModels = $derived.by(() => {
+    const models: string[] = [];
+    for (const [, wrapper] of Object.entries(instancesData)) {
+      if (wrapper && typeof wrapper === "object") {
+        const values = Object.values(wrapper as Record<string, unknown>);
+        if (values.length > 0) {
+          const instance = values[0];
+          if (instance && typeof instance === "object") {
+            const inst = instance as {
+              shardAssignments?: { modelId?: string };
+            };
+            const modelId = inst.shardAssignments?.modelId;
+            if (modelId && !models.includes(modelId)) {
+              models.push(modelId);
+            }
+          }
+        }
+      }
+    }
+    return models;
+  });
+
+  function estimateParamSize(modelId: string): number {
+    const match = modelId.match(/(\d+(?:\.\d+)?)[Bb]/);
+    return match ? parseFloat(match[1]) : 0;
+  }
+
+  const modelsBySize = $derived(
+    [...runningModels].sort(
+      (a, b) => estimateParamSize(b) - estimateParamSize(a),
+    ),
+  );
+
+  const defaultTiers = $derived.by(() => {
+    const n = modelsBySize.length;
+    if (n === 0)
+      return {
+        opus: "your-model-id",
+        sonnet: "your-model-id",
+        haiku: "your-model-id",
+      };
+    if (n === 1)
+      return {
+        opus: modelsBySize[0],
+        sonnet: modelsBySize[0],
+        haiku: modelsBySize[0],
+      };
+    if (n === 2)
+      return {
+        opus: modelsBySize[0],
+        sonnet: modelsBySize[1],
+        haiku: modelsBySize[1],
+      };
+    return {
+      opus: modelsBySize[0],
+      sonnet: modelsBySize[Math.floor(n / 2)],
+      haiku: modelsBySize[n - 1],
+    };
+  });
+
+  let opusModel = $state("");
+  let sonnetModel = $state("");
+  let haikuModel = $state("");
+
+  $effect(() => {
+    opusModel = defaultTiers.opus;
+    sonnetModel = defaultTiers.sonnet;
+    haikuModel = defaultTiers.haiku;
+  });
+
+  let codexModel = $state("");
+  let codexMcpPath = $state("/Users/username");
+  let openClawModel = $state("");
+  $effect(() => {
+    const def = modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id";
+    codexModel = def;
+    openClawModel = def;
+  });
+
+  const claudeShellCommand = $derived(
+    [
+      `ANTHROPIC_BASE_URL=${apiUrl} \\`,
+      `ANTHROPIC_API_KEY=x \\`,
+      `ANTHROPIC_DEFAULT_OPUS_MODEL=${opusModel} \\`,
+      `ANTHROPIC_DEFAULT_SONNET_MODEL=${sonnetModel} \\`,
+      `ANTHROPIC_DEFAULT_HAIKU_MODEL=${haikuModel} \\`,
+      `API_TIMEOUT_MS=3000000 \\`,
+      `CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1 \\`,
+      `claude`,
+    ].join("\n"),
+  );
+
+  const claudeSettingsJson = $derived(
+    JSON.stringify(
+      {
+        env: {
+          ANTHROPIC_BASE_URL: apiUrl,
+          ANTHROPIC_API_KEY: "x",
+          ANTHROPIC_DEFAULT_OPUS_MODEL: opusModel,
+          ANTHROPIC_DEFAULT_SONNET_MODEL: sonnetModel,
+          ANTHROPIC_DEFAULT_HAIKU_MODEL: haikuModel,
+          API_TIMEOUT_MS: "3000000",
+          CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC: "1",
+        },
+      },
+      null,
+      2,
+    ),
+  );
+
+  const openCodeConfig = $derived.by(() => {
+    const models: Record<string, Record<string, unknown>> = {};
+    for (const modelId of runningModels) {
+      const caps = modelCapabilities[modelId] || [];
+      const ctxLen = modelContextLengths[modelId] || 0;
+      const entry: Record<string, unknown> = { name: modelId };
+      if (ctxLen > 0) {
+        entry.limit = { context: ctxLen, output: Math.min(ctxLen, 16384) };
+      }
+      if (caps.includes("vision")) {
+        entry.modalities = { input: ["text", "image"], output: ["text"] };
+      }
+      models[modelId] = entry;
+    }
+    if (Object.keys(models).length === 0) {
+      models["your-model-id"] = { name: "your-model-name" };
+    }
+    const firstModel =
+      runningModels.length > 0 ? runningModels[0] : "your-model-id";
+    return JSON.stringify(
+      {
+        $schema: "https://opencode.ai/config.json",
+        provider: {
+          exo: {
+            npm: "@ai-sdk/openai-compatible",
+            name: "exo",
+            options: {
+              baseURL: `${apiUrl}/v1`,
+              apiKey: "x",
+            },
+            models,
+          },
+        },
+        model: `exo/${firstModel}`,
+      },
+      null,
+      2,
+    );
+  });
+
+  const codexShellCommand = $derived(`EXO_API_KEY=x npx @openai/codex`);
+
+  const codexConfig = $derived(
+    [
+      `model = "${codexModel}"`,
+      `model_provider = "exo"`,
+      ``,
+      `[model_providers.exo]`,
+      `name = "exo"`,
+      `base_url = "${apiUrl}/v1"`,
+      `env_key = "EXO_API_KEY"`,
+      ``,
+      `[mcp_servers.filesystem]`,
+      `command = "npx"`,
+      `args = ["-y", "@modelcontextprotocol/server-filesystem", "${codexMcpPath}"]`,
+    ].join("\n"),
+  );
+
+  const openClawConfig = $derived(
+    JSON.stringify(
+      {
+        gateway: { mode: "local" },
+        models: {
+          providers: {
+            exo: {
+              baseUrl: `${apiUrl}/v1`,
+              apiKey: "x",
+              api: "openai-completions",
+              models: [
+                {
+                  id: openClawModel,
+                  name: "exo local",
+                  input: (modelCapabilities[openClawModel] || []).includes(
+                    "vision",
+                  )
+                    ? ["text", "image"]
+                    : ["text"],
+                },
+              ],
+            },
+          },
+        },
+        agents: {
+          defaults: {
+            model: `exo/${openClawModel}`,
+          },
+        },
+      },
+      null,
+      2,
+    ),
+  );
+
+  const ollamaCommand = $derived(
+    `OLLAMA_HOST=${apiUrl}/ollama ollama run ${modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id"}`,
+  );
+
+  const openWebUiCommand = $derived(
+    [
+      `docker run -d -p 3000:8080 \\`,
+      `  -e OLLAMA_BASE_URL=${apiUrl.replace("localhost", "host.docker.internal")}/ollama \\`,
+      `  -v open-webui:/app/backend/data \\`,
+      `  --name open-webui \\`,
+      `  ghcr.io/open-webui/open-webui:main`,
+    ].join("\n"),
+  );
+
+  const n8nDockerCommand = $derived(
+    [
+      `docker run -d -p 5678:5678 \\`,
+      `  -v n8n_data:/home/node/.n8n \\`,
+      `  --name n8n \\`,
+      `  docker.n8n.io/n8nio/n8n`,
+    ].join("\n"),
+  );
+
+  const n8nCredentialSteps = $derived(
+    [
+      `1. Go to Credentials → Add Credential → search "OpenAI API"`,
+      `2. Set API Key to: x`,
+      `3. Set Base URL to: ${apiUrl.replace("127.0.0.1", "host.docker.internal").replace("localhost", "host.docker.internal")}/v1`,
+      `4. Save the credential`,
+    ].join("\n"),
+  );
+
+  const n8nWorkflowSteps = $derived(
+    [
+      `1. Create a new workflow → "Start from Scratch"`,
+      `2. Add an "AI Agent" or "Basic LLM Chain" node`,
+      `3. Inside it, add an "OpenAI Chat Model" sub-node`,
+      `4. Select the OpenAI credential you just created`,
+      `5. Set Model to "From list" and pick your model (e.g. ${modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id"})`,
+      `6. Optionally toggle "Use Responses API", add Built-in Tools, or click "Add Option" for sampling settings`,
+      `7. Connect a "Chat Trigger" node for interactive chat`,
+      `8. On the Chat Trigger, enable "Allow File Uploads" for vision`,
+    ].join("\n"),
+  );
+
+  const firefoxConfig = $derived(
+    [
+      `1. Open about:config in Firefox`,
+      `2. Set browser.ml.chat.enabled to true`,
+      `3. Set browser.ml.chat.hideLocalhost to false`,
+      `4. Set browser.ml.chat.provider to: ${apiUrl}/`,
+    ].join("\n"),
+  );
+
+  const tabs = [
+    "Claude Code",
+    "OpenCode",
+    "Codex",
+    "OpenClaw",
+    "Open WebUI",
+    "n8n",
+    "Firefox",
+  ] as const;
+  type Tab = (typeof tabs)[number];
+  const stored = browser ? localStorage.getItem("exo-integrations-tab") : null;
+  let activeTab = $state<Tab>(
+    stored && tabs.includes(stored as Tab) ? (stored as Tab) : "Claude Code",
+  );
+  $effect(() => {
+    if (browser) localStorage.setItem("exo-integrations-tab", activeTab);
+  });
+
+  const selectClass =
+    "bg-black/30 border border-exo-light-gray/20 rounded px-2 py-1.5 text-white font-mono text-xs focus:border-exo-yellow/50 focus:outline-none appearance-none cursor-pointer";
+
+  onMount(async () => {
+    refreshState();
+    try {
+      const resp = await fetch("/v1/models");
+      const data = (await resp.json()) as {
+        data: { id: string; capabilities: string[]; context_length: number }[];
+      };
+      const caps: Record<string, string[]> = {};
+      const ctxs: Record<string, number> = {};
+      for (const model of data.data) {
+        caps[model.id] = model.capabilities || [];
+        if (model.context_length > 0) ctxs[model.id] = model.context_length;
+      }
+      modelCapabilities = caps;
+      modelContextLengths = ctxs;
+    } catch {
+      /* ignore */
+    }
+  });
+</script>
+
+<div class="min-h-screen bg-exo-dark-gray flex flex-col">
+  <HeaderNav showHome={true} />
+
+  <main
+    class="flex-1 max-w-3xl mx-auto w-full px-4 md:px-6 py-8"
+    in:fade={{ duration: 200 }}
+  >
+    <div class="mb-8">
+      <h1
+        class="text-white text-xl md:text-2xl font-semibold tracking-wide mb-2"
+      >
+        Integrations
+      </h1>
+      <p class="text-exo-light-gray/60 text-sm">
+        Connect external tools to your exo cluster.
+      </p>
+    </div>
+
+    <!-- Status -->
+    <div class="mb-8">
+      <span class="text-exo-light-gray/70 text-xs uppercase tracking-wider"
+        >API Endpoint</span
+      >
+      <span class="text-white font-mono text-sm ml-2">{apiUrl}</span>
+      {#if runningModels.length > 0}
+        <div class="text-exo-light-gray/50 text-xs mt-2">
+          Running model{runningModels.length > 1 ? "s" : ""}:
+          <ul class="mt-1 space-y-0.5 list-none">
+            {#each runningModels as model}
+              <li class="text-exo-yellow font-mono">{model}</li>
+            {/each}
+          </ul>
+        </div>
+      {:else}
+        <p class="text-exo-light-gray/40 text-xs mt-2 italic">
+          No models currently running
+        </p>
+      {/if}
+    </div>
+
+    <!-- API Endpoints -->
+    <div class="mb-8">
+      <div
+        class="flex flex-col sm:flex-row gap-3 text-xs font-mono text-exo-light-gray/70"
+      >
+        <div
+          class="flex-1 bg-black/20 border border-exo-light-gray/10 rounded px-3 py-2"
+        >
+          <span class="text-exo-light-gray/40 text-[10px] uppercase block mb-1"
+            >OpenAI-compatible</span
+          >
+          <span class="text-white/80">{apiUrl}/v1</span>
+        </div>
+        <div
+          class="flex-1 bg-black/20 border border-exo-light-gray/10 rounded px-3 py-2"
+        >
+          <span class="text-exo-light-gray/40 text-[10px] uppercase block mb-1"
+            >Claude-compatible</span
+          >
+          <span class="text-white/80">{apiUrl}</span>
+        </div>
+        <div
+          class="flex-1 bg-black/20 border border-exo-light-gray/10 rounded px-3 py-2"
+        >
+          <span class="text-exo-light-gray/40 text-[10px] uppercase block mb-1"
+            >Ollama-compatible</span
+          >
+          <span class="text-white/80">{apiUrl}/ollama</span>
+        </div>
+      </div>
+    </div>
+
+    <!-- Tabs -->
+    <div
+      class="flex flex-wrap gap-2 mb-6 border-b border-exo-light-gray/10 pb-3"
+    >
+      {#each tabs as tab}
+        <button
+          onclick={() => (activeTab = tab)}
+          class="px-3 py-1.5 text-xs rounded-md transition-all cursor-pointer
+            {activeTab === tab
+            ? 'bg-exo-yellow/15 text-exo-yellow border border-exo-yellow/30'
+            : 'text-exo-light-gray/60 hover:text-white/80 border border-transparent hover:border-exo-light-gray/20'}"
+        >
+          {tab}
+        </button>
+      {/each}
+    </div>
+
+    <!-- Tab Content -->
+    <div class="space-y-4">
+      {#if activeTab === "Claude Code"}
+        {#if runningModels.length > 1}
+          <div class="grid grid-cols-3 gap-3 text-xs">
+            {#each [{ label: "Opus", bind: () => opusModel, set: (v: string) => (opusModel = v) }, { label: "Sonnet", bind: () => sonnetModel, set: (v: string) => (sonnetModel = v) }, { label: "Haiku", bind: () => haikuModel, set: (v: string) => (haikuModel = v) }] as tier}
+              <div>
+                <span
+                  class="text-exo-light-gray/50 text-[10px] uppercase tracking-wider block mb-1"
+                  >{tier.label}</span
+                >
+                <select
+                  value={tier.bind()}
+                  onchange={(e) =>
+                    tier.set((e.target as HTMLSelectElement).value)}
+                  class="w-full {selectClass}"
+                >
+                  {#each runningModels as model}
+                    <option value={model}>{model.split("/").pop()}</option>
+                  {/each}
+                </select>
+              </div>
+            {/each}
+          </div>
+        {/if}
+        <IntegrationCard
+          title="Shell Command"
+          subtitle="Run in terminal"
+          description="Launch Claude Code with exo as the backend. Paste this into your terminal."
+          config={claudeShellCommand}
+          language="bash"
+        />
+        <IntegrationCard
+          title="Settings File"
+          subtitle="~/.claude/settings.json"
+          description="Or add this to your Claude Code settings for persistent configuration."
+          config={claudeSettingsJson}
+        />
+      {:else if activeTab === "OpenCode"}
+        <IntegrationCard
+          title="Config File"
+          subtitle="opencode.json"
+          description="Add this to your project root or ~/.config/opencode/opencode.json for global config. Vision models automatically get image input modality."
+          config={openCodeConfig}
+        />
+      {:else if activeTab === "Codex"}
+        <div class="flex gap-3 text-xs">
+          {#if runningModels.length > 1}
+            <div>
+              <span
+                class="text-exo-light-gray/50 text-[10px] uppercase tracking-wider block mb-1"
+                >Model</span
+              >
+              <select bind:value={codexModel} class={selectClass}>
+                {#each runningModels as model}
+                  <option value={model}>{model.split("/").pop()}</option>
+                {/each}
+              </select>
+            </div>
+          {/if}
+          <div class="flex-1">
+            <span
+              class="text-exo-light-gray/50 text-[10px] uppercase tracking-wider block mb-1"
+              >MCP Filesystem Path</span
+            >
+            <input
+              type="text"
+              bind:value={codexMcpPath}
+              class="w-full bg-black/30 border border-exo-light-gray/20 rounded px-2 py-1.5 text-white font-mono text-xs focus:border-exo-yellow/50 focus:outline-none"
+            />
+          </div>
+        </div>
+        <IntegrationCard
+          title="Config File"
+          subtitle="~/.codex/config.toml"
+          description="Add this to your Codex CLI config so the model and provider persist."
+          config={codexConfig}
+        />
+        <IntegrationCard
+          title="Shell Command"
+          subtitle="Run in terminal"
+          description="Launch Codex with exo as the backend."
+          config={codexShellCommand}
+          language="bash"
+        />
+      {:else if activeTab === "OpenClaw"}
+        {#if runningModels.length > 1}
+          <div class="text-xs">
+            <span
+              class="text-exo-light-gray/50 text-[10px] uppercase tracking-wider block mb-1"
+              >Model</span
+            >
+            <select bind:value={openClawModel} class={selectClass}>
+              {#each runningModels as model}
+                <option value={model}>{model.split("/").pop()}</option>
+              {/each}
+            </select>
+          </div>
+        {/if}
+        <IntegrationCard
+          title="Config File"
+          subtitle="~/.openclaw/openclaw.json"
+          description="Add this to your OpenClaw config. If you haven't installed OpenClaw yet, run: npm install -g openclaw@latest"
+          config={openClawConfig}
+        />
+        <IntegrationCard
+          title="Setup Commands"
+          subtitle="Run in terminal"
+          description="After saving the config, run these commands to fix metadata and start the gateway."
+          config={`openclaw doctor --fix${(modelCapabilities[openClawModel] || []).includes("vision") ? `\nopenclaw models set-image exo/${openClawModel}` : ""}\nopenclaw gateway &\nopenclaw dashboard`}
+          language="bash"
+        />
+      {:else if activeTab === "Open WebUI"}
+        <IntegrationCard
+          title="1. Start Open WebUI"
+          subtitle="Run in terminal"
+          description="Run this to start Open WebUI."
+          config={openWebUiCommand}
+          language="bash"
+        />
+        <IntegrationCard
+          title="2. Open & Select Model"
+          subtitle="http://localhost:3000"
+          description={`Open http://localhost:3000 in your browser. Select the running model from the dropdown at the top: ${runningModels.length > 0 ? runningModels.join(", ") : "no models running"}`}
+          config={"open http://localhost:3000"}
+          language="bash"
+        />
+        <IntegrationCard
+          title="Ollama CLI"
+          subtitle="Run in terminal"
+          description="Or use the Ollama CLI directly."
+          config={ollamaCommand}
+          language="bash"
+        />
+      {:else if activeTab === "n8n"}
+        <IntegrationCard
+          title="1. Start n8n"
+          subtitle="Run in terminal"
+          description="Start n8n with Docker. If you already have n8n running, skip this step."
+          config={n8nDockerCommand}
+          language="bash"
+        />
+        <IntegrationCard
+          title="2. Open n8n"
+          subtitle="http://localhost:5678"
+          description="Open n8n in your browser. If this is your first time, complete the setup and select 'Start from Scratch' when prompted."
+          config={"open http://localhost:5678"}
+          language="bash"
+        />
+        <IntegrationCard
+          title="3. Add OpenAI Credential"
+          subtitle="n8n UI → Credentials"
+          description="Create an OpenAI credential pointing at your exo cluster."
+          config={n8nCredentialSteps}
+        />
+        <IntegrationCard
+          title="4. Build a Workflow"
+          subtitle="n8n UI → Workflows"
+          description="Create a workflow that uses your exo-powered model."
+          config={n8nWorkflowSteps}
+        />
+      {:else if activeTab === "Firefox"}
+        <IntegrationCard
+          title="Firefox AI Chatbot"
+          subtitle="about:config"
+          description="Use the exo dashboard as Firefox's built-in AI chatbot. Requires Firefox 130+."
+          config={firefoxConfig}
+        />
+      {/if}
+    </div>
+  </main>
+</div>

docs/architecture.md

@@ -81,4 +81,4 @@ Whenever a device produces side effects, it captures those side effects in an `E
 
 ## Purity
 
-A significant goal of the current design is to make data flow explicit. Classes should either represent simple data (`CamelCaseModel`s typically, and `TaggedModel`s for unions) or active `System`s (Erlang `Actor`s), with all transformations of that data being "referentially transparent" - destructure and construct new data, don't mutate in place. We have had varying degrees of success with this, and are still exploring where purity makes sense.
+A significant goal of the current design is to make data flow explicit. Classes should either represent simple data (`FrozenModel`s typically, and `TaggedModel`s for unions) or active `System`s (Erlang `Actor`s), with all transformations of that data being "referentially transparent" - destructure and construct new data, don't mutate in place. We have had varying degrees of success with this, and are still exploring where purity makes sense.

flake.lock

@@ -2,11 +2,11 @@
   "nodes": {
     "crane": {
       "locked": {
-        "lastModified": 1767744144,
-        "narHash": "sha256-9/9ntI0D+HbN4G0TrK3KmHbTvwgswz7p8IEJsWyef8Q=",
+        "lastModified": 1775790182,
+        "narHash": "sha256-pG2RWVQY0Pe+rmmXJx+Jpyi+JcgjWzS18m7fcD1B64Q=",
         "owner": "ipetkov",
         "repo": "crane",
-        "rev": "2fb033290bf6b23f226d4c8b32f7f7a16b043d7e",
+        "rev": "534982f1c41834b101e381b07b1121a4f065a374",
         "type": "github"
       },
       "original": {
@@ -47,11 +47,11 @@
         "rust-analyzer-src": "rust-analyzer-src"
       },
       "locked": {
-        "lastModified": 1768287139,
-        "narHash": "sha256-nsXFt0OzUi6K7dUzzJD5/v9e0Ic+fvclfIW936/43ZM=",
+        "lastModified": 1775807984,
+        "narHash": "sha256-Redoe3D9zGN5I9QPHWL9vfMVQBehY1fKsMiRXQ83X3w=",
         "owner": "nix-community",
         "repo": "fenix",
-        "rev": "a4a3aa956931f90f35453cb519e4545e9ad7f773",
+        "rev": "fcf90c0c4d368b2ca917a7afa6d08e98a397e5fd",
         "type": "github"
       },
       "original": {
@@ -83,11 +83,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1768135262,
-        "narHash": "sha256-PVvu7OqHBGWN16zSi6tEmPwwHQ4rLPU9Plvs8/1TUBY=",
+        "lastModified": 1775087534,
+        "narHash": "sha256-91qqW8lhL7TLwgQWijoGBbiD4t7/q75KTi8NxjVmSmA=",
         "owner": "hercules-ci",
         "repo": "flake-parts",
-        "rev": "80daad04eddbbf5a4d883996a73f3f542fa437ac",
+        "rev": "3107b77cd68437b9a76194f0f7f9c55f2329ca5b",
         "type": "github"
       },
       "original": {
@@ -96,38 +96,42 @@
         "type": "github"
       }
     },
+    "nixglhost": {
+      "inputs": {
+        "nixpkgs": [
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1732211616,
+        "narHash": "sha256-QZCKJoypcwgS3tDNSWMjlxEBZtOYPW3eXV24rMzKsac=",
+        "owner": "numtide",
+        "repo": "nix-gl-host",
+        "rev": "5269b233f83880a0b433eafe026f0bc0d8f1a4a9",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "nix-gl-host",
+        "type": "github"
+      }
+    },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1768127708,
-        "narHash": "sha256-1Sm77VfZh3mU0F5OqKABNLWxOuDeHIlcFjsXeeiPazs=",
+        "lastModified": 1775595990,
+        "narHash": "sha256-OEf7YqhF9IjJFYZJyuhAypgU+VsRB5lD4DuiMws5Ltc=",
         "owner": "NixOS",
         "repo": "nixpkgs",
-        "rev": "ffbc9f8cbaacfb331b6017d5a5abb21a492c9a38",
+        "rev": "4e92bbcdb030f3b4782be4751dc08e6b6cb6ccf2",
         "type": "github"
       },
       "original": {
         "owner": "NixOS",
-        "ref": "nixos-unstable",
+        "ref": "nixos-25.11",
         "repo": "nixpkgs",
         "type": "github"
       }
     },
-    "nixpkgs-swift": {
-      "locked": {
-        "lastModified": 1761672384,
-        "narHash": "sha256-o9KF3DJL7g7iYMZq9SWgfS1BFlNbsm6xplRjVlOCkXI=",
-        "owner": "NixOS",
-        "repo": "nixpkgs",
-        "rev": "08dacfca559e1d7da38f3cf05f1f45ee9bfd213c",
-        "type": "github"
-      },
-      "original": {
-        "owner": "NixOS",
-        "repo": "nixpkgs",
-        "rev": "08dacfca559e1d7da38f3cf05f1f45ee9bfd213c",
-        "type": "github"
-      }
-    },
     "purescript-overlay": {
       "inputs": {
         "flake-compat": "flake-compat",
@@ -164,11 +168,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1763662255,
-        "narHash": "sha256-4bocaOyLa3AfiS8KrWjZQYu+IAta05u3gYZzZ6zXbT0=",
+        "lastModified": 1773870109,
+        "narHash": "sha256-ZoTdqZP03DcdoyxvpFHCAek4bkPUTUPUF3oCCgc3dP4=",
         "owner": "pyproject-nix",
         "repo": "build-system-pkgs",
-        "rev": "042904167604c681a090c07eb6967b4dd4dae88c",
+        "rev": "b6e74f433b02fa4b8a7965ee24680f4867e2926f",
         "type": "github"
       },
       "original": {
@@ -184,11 +188,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1764134915,
-        "narHash": "sha256-xaKvtPx6YAnA3HQVp5LwyYG1MaN4LLehpQI8xEdBvBY=",
+        "lastModified": 1775439158,
+        "narHash": "sha256-NHY9SJNU019n+8NCabBDtmuzRFeE2gZlYKHowp9bV24=",
         "owner": "pyproject-nix",
         "repo": "pyproject.nix",
-        "rev": "2c8df1383b32e5443c921f61224b198a2282a657",
+        "rev": "fb6b728260f3f32761367e9fd1e1a25b4245bcd0",
         "type": "github"
       },
       "original": {
@@ -203,8 +207,8 @@
         "dream2nix": "dream2nix",
         "fenix": "fenix",
         "flake-parts": "flake-parts",
+        "nixglhost": "nixglhost",
         "nixpkgs": "nixpkgs",
-        "nixpkgs-swift": "nixpkgs-swift",
         "pyproject-build-systems": "pyproject-build-systems",
         "pyproject-nix": "pyproject-nix",
         "treefmt-nix": "treefmt-nix",
@@ -214,11 +218,11 @@
     "rust-analyzer-src": {
       "flake": false,
       "locked": {
-        "lastModified": 1768224240,
-        "narHash": "sha256-Pp1dDrXKPBUJReZnnDElFyHYn67XTd48zRhToheLjtk=",
+        "lastModified": 1775745684,
+        "narHash": "sha256-8MbfLwd60FNa8dRFkjE+G3TT/x21G3Rsplm1bMBQUtU=",
         "owner": "rust-lang",
         "repo": "rust-analyzer",
-        "rev": "725349602e525df37f377701e001fe8aab807878",
+        "rev": "64ddb549bc9a70d011328746fa46a8883f937b6b",
         "type": "github"
       },
       "original": {
@@ -257,11 +261,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1768158989,
-        "narHash": "sha256-67vyT1+xClLldnumAzCTBvU0jLZ1YBcf4vANRWP3+Ak=",
+        "lastModified": 1775636079,
+        "narHash": "sha256-pc20NRoMdiar8oPQceQT47UUZMBTiMdUuWrYu2obUP0=",
         "owner": "numtide",
         "repo": "treefmt-nix",
-        "rev": "e96d59dff5c0d7fddb9d113ba108f03c3ef99eca",
+        "rev": "790751ff7fd3801feeaf96d7dc416a8d581265ba",
         "type": "github"
       },
       "original": {
@@ -280,11 +284,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1767701098,
-        "narHash": "sha256-CJhKZnWb3gumR9oTRjFvCg/6lYTGbZRU7xtvcyWIRwU=",
+        "lastModified": 1775706324,
+        "narHash": "sha256-BTb4sydzX2B5/oNbvCdQFeSbk97xEnbb8bk84CiKCOs=",
         "owner": "pyproject-nix",
         "repo": "uv2nix",
-        "rev": "9d357f0d2ce6f5f35ec7959d7e704452352eb4da",
+        "rev": "5707df99097375896a3dda811d492a2fabe63500",
         "type": "github"
       },
       "original": {

flake.nix

@@ -2,7 +2,7 @@
   description = "The development environment for Exo";
 
   inputs = {
-    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
+    nixpkgs.url = "github:NixOS/nixpkgs/nixos-25.11";
 
     flake-parts = {
       url = "github:hercules-ci/flake-parts";
@@ -46,17 +46,18 @@
       inputs.nixpkgs.follows = "nixpkgs";
     };
 
-    # Pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
-    nixpkgs-swift.url = "github:NixOS/nixpkgs/08dacfca559e1d7da38f3cf05f1f45ee9bfd213c";
+    nixglhost = {
+      url = "github:numtide/nix-gl-host";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
   };
 
   nixConfig = {
-    extra-trusted-public-keys = "exo.cachix.org-1:okq7hl624TBeAR3kV+g39dUFSiaZgLRkLsFBCuJ2NZI=";
-    extra-substituters = "https://exo.cachix.org";
+    extra-trusted-public-keys = "exo.cachix.org-1:okq7hl624TBeAR3kV+g39dUFSiaZgLRkLsFBCuJ2NZI= cache.nixos-cuda.org:74DUi4Ye579gUqzH4ziL9IyiJBlDpMRn9MBN8oNan9M=";
+    extra-substituters = "https://exo.cachix.org https://cache.nixos-cuda.org";
   };
 
-  outputs =
-    inputs:
+  outputs = inputs:
     inputs.flake-parts.lib.mkFlake { inherit inputs; } {
       systems = [
         "x86_64-linux"
@@ -71,21 +72,45 @@
         ./python/parts.nix
       ];
 
+      debug = true; # Enable options autocompletion
+
       perSystem =
-        { config, self', inputs', pkgs, lib, system, ... }:
+        { config
+        , self'
+        , pkgs
+        , lib
+        , system
+        , ...
+        }:
         let
-          # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
-          pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
-        in
-        {
-          # Allow unfree for metal-toolchain (needed for Darwin Metal packages)
-          _module.args.pkgs = import inputs.nixpkgs {
+          pkgsArgs = {
             inherit system;
             config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain";
             overlays = [
+              inputs.nixglhost.overlays.default
               (import ./nix/apple-sdk-overlay.nix)
+              (final: _: {
+                macmon = final.rustPlatform.buildRustPackage {
+                  pname = "macmon";
+                  version = "git";
+                  src = final.fetchFromGitHub {
+                    owner = "vladkens";
+                    repo = "macmon";
+                    rev = "a1cd06b6cc0d5e61db24fd8832e74cd992097a7d";
+                    hash = "sha256-wcq4PUXK44XfUKOZKl32u8LpOxXpSbUUfItQGwS2Zso=";
+                  };
+                  cargoHash = "sha256-Epj3L+db1flGNK5y6yfSig8piEiXTz15lPo/FNkqlkA=";
+                };
+              })
             ];
           };
+        in
+        {
+          # Allow unfree for metal-toolchain (needed for Darwin Metal packages)
+          _module.args = {
+            pkgs = import inputs.nixpkgs pkgsArgs;
+            unfreePkgs = import inputs.nixpkgs (pkgsArgs // { config.allowUnfree = true; });
+          };
           treefmt = {
             projectRootFile = "flake.nix";
             programs = {
@@ -105,74 +130,63 @@
               };
               swift-format = {
                 enable = true;
-                package = pkgsSwift.swiftPackages.swift-format;
+                package = pkgs.swiftPackages.swift-format;
               };
               shfmt.enable = true;
               taplo.enable = true;
             };
           };
 
-          packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
-            let
-              uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
-              mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
-              uvLockMlxVersion = mlxPackage.version;
-              uvLockMlxRev = builtins.elemAt (builtins.split "#" mlxPackage.source.git) 2;
-            in
+          packages =
             {
-              metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
-              mlx = pkgs.callPackage ./nix/mlx.nix {
-                inherit (self'.packages) metal-toolchain;
-                inherit uvLockMlxVersion uvLockMlxRev;
-              };
               default = self'.packages.exo;
             }
-          );
+            // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
+              metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
+            };
 
-          devShells.default = with pkgs; pkgs.mkShell {
-            inputsFrom = [ self'.checks.cargo-build ];
+          devShells.default = with pkgs;
+            pkgs.mkShell {
+              inputsFrom = [ self'.checks.cargo-build ];
 
-            packages =
-              [
-                # FORMATTING
-                config.treefmt.build.wrapper
+              packages =
+                [
+                  # FORMATTING
+                  config.treefmt.build.wrapper
 
-                # PYTHON
-                python313
-                uv
-                ruff
-                basedpyright
+                  # PYTHON
+                  self'.packages.editableVenv
+                  uv
 
-                # RUST
-                config.rust.toolchain
-                maturin
+                  # RUST
+                  config.rust.toolchain
+                  maturin
 
-                # NIX
-                nixpkgs-fmt
+                  # NIX
+                  nixd
+                  nixpkgs-fmt
 
-                # SVELTE
-                nodejs
+                  # SVELTE
+                  nodejs
 
-                # MISC
-                just
-                jq
-              ]
-              ++ lib.optionals stdenv.isLinux [
-                unixtools.ifconfig
-              ]
-              ++ lib.optionals stdenv.isDarwin [
-                macmon
-              ];
+                  # MISC
+                  just
+                  jq
+                ]
+                ++ lib.optionals stdenv.isDarwin [
+                  macmon
+                  self'.packages.metal-toolchain
+                ];
 
-            OPENSSL_NO_VENDOR = "1";
+              OPENSSL_NO_VENDOR = "1";
 
-            shellHook = ''
-              export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
-              ${lib.optionalString stdenv.isLinux ''
-                export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
-              ''}
-            '';
-          };
+              shellHook = ''
+                export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
+                ${lib.optionalString stdenv.isLinux ''
+                  export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
+                ''}
+              '';
+            };
         };
     };
 }

justfile

@@ -1,5 +1,8 @@
 export NIX_CONFIG := "extra-experimental-features = nix-command flakes"
 
+default: lint fmt
+all: lint fmt check
+
 fmt:
     treefmt || nix fmt
 
@@ -19,17 +22,23 @@ sync-clean:
     uv sync --all-packages --force-reinstall --no-cache
 
 rust-rebuild:
-    cargo run --bin stub_gen
+    PYO3_PYTHON="$(uv run python -c 'import sys; print(sys.executable)')" cargo run --bin stub_gen
     uv sync --reinstall-package exo_pyo3_bindings
 
 build-dashboard:
     #!/usr/bin/env bash
-    cd dashboard
+    pushd dashboard
     npm install
     npm run build
+    popd
 
-package:
+package: build-dashboard
     uv run pyinstaller packaging/pyinstaller/exo.spec
+    rm -rf build
+
+build-app: rust-rebuild sync-clean package
+    xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
+    @echo "\nBuild complete. Run with:\n  open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
 
 clean:
     rm -rf **/__pycache__

mlx_tinygrad_interop/README.md

@@ -0,0 +1,251 @@
+# MLX Tinygrad Interop
+
+Private code for benchmarking MLX `<->` tinygrad tensor conversions.
+
+There is also now a lightweight baseline bridge module in
+`mlx_tinygrad_interop/lib/tensor_bridge.py`, modelled after vLLM Metal's
+`tensor_bridge.py` shape:
+
+- `tinygrad_to_mlx(...)` is implemented
+- `mlx_to_tinygrad(...)` is intentionally stubbed there
+
+That module is meant as a simple public-shaped bridge baseline, separate from
+the lower-level benchmark and lease-pool experiments.
+
+Reusable runtime code now lives under `mlx_tinygrad_interop/lib/`.
+Benchmarks, stress harnesses, and tests stay in the top-level
+`mlx_tinygrad_interop/` package.
+
+There is also a separate benchmark file for the route that relies on existing
+tinygrad <-> PyTorch and MLX <-> PyTorch interop instead of new framework
+patches:
+
+- `mlx_tinygrad_interop/bench_torch_route.py`
+
+That route requires `torch` in the normal macOS `exo` environment.
+The currently working fully pre-existing route is the CPU-intermediate one.
+Using an intermediate PyTorch `mps` tensor caused the documented
+`Tensor.from_blob(..., device="METAL")` path to fail at runtime on `e16`, so
+the benchmark defaults to `--torch-device cpu`.
+
+## Workflow
+
+Use the repo devshell and top-level dependency graph. Do not install ad-hoc
+build dependencies or patch around them with one-off environment setups.
+
+1. Change code locally.
+2. Push the `mlx` and `tinygrad` fork changes.
+3. In local `exo`, enter the devshell with `nix develop`.
+4. Refresh `uv.lock` against the new fork heads with:
+   `uv lock --upgrade-package mlx --refresh-package mlx --upgrade-package tinygrad --refresh-package tinygrad`
+5. Commit and push the updated `exo` branch.
+6. On the remote Mac, pull the updated repos.
+7. Enter the devshell with `nix develop`.
+8. Refresh the environment with `uv sync`.
+9. Run tests or benchmarks with `uv run ...`.
+
+Plain `uv lock` was not enough to move these git-based dependency SHAs during
+testing, and `--upgrade-package` alone still left one stale git revision in a
+later pass. The working command was the explicit `--upgrade-package` plus
+`--refresh-package` form above.
+
+## Benchmark
+
+The current benchmark keeps source tensor construction and explicit pre-sync
+outside the timed loop, but many rows still include per-call helper, binding,
+owner pinning, and wrapper-construction overhead.
+
+Current benchmark CSV output reports:
+
+- `avg_us`
+- `stddev_us`
+
+Older notes in this file that mention min/median refer to earlier runs before
+the reporting format was changed.
+
+- Inputs are assumed to already be synchronized.
+- Inputs are assumed to already be allocated.
+- Inputs are assumed to already be materialized / realized.
+- Setup stays outside the timed loop.
+- The current unsafe helper bridge is asymmetric:
+  - `MLX -> tinygrad` adopts an existing `MTLBuffer*`
+  - `tinygrad -> MLX` rebuilds an MLX array from a raw pointer
+- Newer `MLX -> tinygrad` rows also cover:
+  - a single MLX-side entrypoint that calls into tinygrad without the exporter
+    dict round-trip
+  - a rebindable tinygrad slot that reuses one wrapper and rebinds the
+    borrowed `MTLBuffer*`
+  - a small ring of such slots
+  - a keyed lease pool that owns `mx.eval(...)` on acquire and explicit lease
+    release on the tinygrad side
+  - `*_then_use_sum` rows that immediately consume the converted tensor through
+    a realized tinygrad reduction
+- `mx.array(memoryview(...))` is a native copy path in current MLX, not an
+  aliasing import path.
+- Rebindable slots enforce fixed shape/dtype contracts on rebind.
+- Do not rebind a slot until all work derived from its previous contents has
+  been realized and synchronized.
+- The practical `MLX -> tinygrad` path is now a lease-managed pool keyed by
+  `(shape, dtype, byte_offset)`, not a bare mutable borrower.
+- There is now also a copy-based `MLX -> tinygrad` pool family:
+  - `MlxToTinygradCopyLeasePool`
+  - `MlxToTinygradCopyLeasePools`
+  - these reuse tinygrad-owned destination tensors and copy MLX bytes into
+    them instead of aliasing a foreign `MTLBuffer*`
+- The raw lease path remains intentionally unsafe:
+  - `lease.tensor` is not a snapshot
+  - if that raw tensor escapes beyond the lease and the slot is reused, it can
+    observe new contents
+- The preferred production-shaped API is the scoped callback form:
+  - `pool.run_with_mlx_tensor(array, fn=...)`
+  - `pools.run_with_mlx_tensor(array, tg_dtype=..., fn=...)`
+  - these scope acquire/use/release together and only allow independently
+    realized outputs to escape
+  - they reject returning the borrowed tensor directly
+  - they reject returning alias views of the borrowed slot
+  - they reject leaked live tensors whose graphs still depend on the borrowed
+    tensor
+  - the callback still must not stash the raw borrowed tensor object itself;
+    that remains a contract rule rather than something the current runtime can
+    prove mechanically
+- Safe scoped release uses `Device["METAL"].synchronize()` again. The narrower
+  callback-local command-buffer wait experiment was not concurrency-safe enough
+  to keep as the default runtime behavior.
+- Alias and copy pool registries are bounded keyed caches with `max_pools`, so
+  variable-shape inference can be bucketed without unbounded registry growth.
+- This fast path is only valid for same-process, same-address-space Apple
+  Silicon unified-memory handoff. It does not cross process or machine
+  boundaries, and it does not remove any later Metal/host -> CUDA transfer.
+
+Test command:
+
+```bash
+uv run python -m unittest mlx_tinygrad_interop.test_interop mlx_tinygrad_interop.test_handoff
+```
+
+Stress command:
+
+```bash
+uv run python mlx_tinygrad_interop/stress_interop.py --cases 64 --soak-iterations 512
+```
+
+The stress suite now also reports native memory signals:
+
+- `mx.get_active_memory()`
+- `mx.get_cache_memory()`
+- `mx.get_peak_memory()`
+- process `ru_maxrss`
+
+It also now checks:
+
+- more complex movement / broadcast / reduction / matmul chains
+- roundtrip `MLX -> tinygrad -> MLX` correctness after those chains
+- bounded alias/copy pool-count behavior during soak runs
+
+For `matmul_lastdim`, the stress harness uses `np.einsum(...)` for the NumPy
+baseline instead of NumPy `@`. On the current macOS validation host, a valid
+contiguous float32 matmul case through `@` returned an incorrect all-zero
+result while MLX, tinygrad, and `np.einsum` agreed on the nonzero output.
+
+Float32 stress comparisons also allow a small `rtol=5e-5, atol=1e-5` tolerance
+so mixed matmul/reduction chains are not failed for a few-ulps backend
+accumulation-order drift.
+
+Raw conversions are still checked against NumPy values directly, but the
+downstream op-chain checks now use the native destination-framework baseline:
+
+- `MLX -> tinygrad` op chains are compared to native tinygrad results
+- `tinygrad -> MLX` op chains are compared to native MLX results
+
+That avoids treating real framework semantic differences, such as integer
+promotion behavior, as interop failures.
+
+Example:
+
+```bash
+uv run python mlx_tinygrad_interop/bench_raw_conversion.py --dtype float32 --sizes 256,512,1024,2048,4096,7168
+```
+
+Validated remote command on `e16`:
+
+```bash
+uv run python mlx_tinygrad_interop/bench_raw_conversion.py --dtype float32 --sizes 7168 --warmup 64 --samples 7 --min-batch-us 1000
+```
+
+Observed `7168`-byte results on that run:
+
+- `unsafe_helper_bridge`
+  - `mlx_to_tinygrad`: `21.202 us` min, `21.532 us` median
+  - `tinygrad_to_mlx`: `28.109 us` min, `28.372 us` median
+- `single_entry_bridge`
+  - `mlx_to_tinygrad`: `21.388 us` min, `21.542 us` median
+- `fresh_wrapper_then_use_sum`
+  - `mlx_to_tinygrad`: `601.812 us` min, `611.458 us` median
+- `rebindable_slot_bridge`
+  - `mlx_to_tinygrad`: `1.505 us` min, `1.542 us` median
+- `rebindable_slot_then_use_sum`
+  - `mlx_to_tinygrad`: `579.583 us` min, `581.833 us` median
+- `borrower_ring4_bridge`
+  - `mlx_to_tinygrad`: `1.531 us` min, `1.573 us` median
+- `borrower_ring4_then_use_sum`
+  - `mlx_to_tinygrad`: `577.730 us` min, `581.000 us` median
+- `unsafe_helper_legacy`
+  - `mlx_to_tinygrad`: `31.938 us` min, `32.214 us` median
+- `unsafe_helper_maybe_copy`
+  - `tinygrad_to_mlx`: `28.153 us` min, `28.277 us` median
+- `memoryview_copy`
+  - `mlx_to_tinygrad`: `35.191 us` min, `35.668 us` median
+  - `tinygrad_to_mlx`: `2.596 us` min, `2.662 us` median
+- `numpy_baseline`
+  - `mlx_to_tinygrad`: `272.323 us` min, `275.104 us` median
+  - `tinygrad_to_mlx`: `12.817 us` min, `13.005 us` median
+
+Later remote microbench runs showed the split more clearly:
+
+- `MLX -> tinygrad single_entry_bridge` barely changes the fresh-wrapper cost,
+  so exporter dict marshalling was never the main problem.
+- `MLX -> tinygrad` is dominated by tinygrad import / wrapper construction when
+  a fresh tensor is created each time.
+- The rebindable tinygrad slot drops `MLX -> tinygrad` to about `1.5 us`, and
+  a ring of four slots stays at essentially the same latency. That means
+  wrapper reuse, not exporter marshalling, is the decisive optimization on this
+  host.
+- After adding hardened shape/dtype contract checks, a spot-check at `7168`
+  bytes moved those rows to about `2.46 us` for the single slot and `2.55 us`
+  for the ring. That is still comfortably inside the target latency range.
+- The strict alias-only `tinygrad -> MLX` helper succeeds on `e16`; its timing
+  is effectively the same as the maybe-copy helper on that host.
+- `tinygrad -> MLX` is dominated by tinygrad export in the unsafe helper path.
+- `tinygrad -> MLX memoryview_copy` is already the practical low-latency path
+  for small tensors.
+- Offsetted MLX slices now export both logical bytes and backing-buffer bytes,
+  and a nonzero-offset slice was validated successfully into tinygrad.
+- The randomized stress suite also caught and fixed the zero-offset variant of
+  that problem: oversized backing buffers now import through a logical tinygrad
+  buffer view instead of reshaping the entire backing allocation.
+- The same stress suite also found that MLX backing-buffer capacity is a raw
+  byte count, not necessarily a dtype-aligned element count. The fast tinygrad
+  path now handles that with byte-level bounds checks and `ceildiv`.
+- The rebindable slot returns the same tinygrad `Tensor` object rebound to new
+  Metal storage, so it is narrower than an ordinary "new tensor each call"
+  conversion helper.
+- The current ring rows are still benchmark primitives, not a production lease
+  API. If this path is used in the real disaggregated MLX/tinygrad runtime, it
+  should use the scoped pool callback API rather than a raw escaping lease
+  tensor wherever possible.
+- Safe lease release now clears the slot's pinned MLX owner reference after the
+  Metal barrier. Unsafe `synchronize_on_release=False` use is still the
+  caller's responsibility.
+- The `*_then_use_sum` rows are dominated by the realized tinygrad reduction.
+  They should be read as end-to-end "convert then immediately consume" probes.
+  They still show the same relative story: slot/ring rebinding saves about
+  `20-25 us` versus the fresh-wrapper path at `7 kB`.
+
+## Current Scope
+
+- Private / unsafe helpers only.
+- Metal / unified-memory path only.
+- Dense contiguous tensors only.
+- Same-dtype conversions only.
+- Current exporter/importer microbenchmarks are intended to separate helper
+  overhead from end-to-end bridge cost.

mlx_tinygrad_interop/__init__.py

@@ -0,0 +1,23 @@
+"""Private MLX <-> tinygrad interop experiments, benchmarks, and handoff helpers."""
+
+from mlx_tinygrad_interop.lib import (
+  MlxToTinygradLease,
+  MlxToTinygradLeaseKey,
+  MlxToTinygradLeasePool,
+  MlxToTinygradLeasePools,
+  mlx_to_tinygrad,
+  sync_mlx,
+  sync_tinygrad,
+  tinygrad_to_mlx,
+)
+
+__all__ = [
+  "MlxToTinygradLease",
+  "MlxToTinygradLeaseKey",
+  "MlxToTinygradLeasePool",
+  "MlxToTinygradLeasePools",
+  "mlx_to_tinygrad",
+  "sync_mlx",
+  "sync_tinygrad",
+  "tinygrad_to_mlx",
+]

mlx_tinygrad_interop/bench_raw_conversion.py

@@ -0,0 +1,443 @@
+import argparse
+import gc
+import platform
+import statistics
+import sys
+import time
+from typing import Any, Callable, cast
+
+import mlx.core as mx
+import numpy as np
+from tinygrad import Device, Tensor, dtypes
+from tinygrad.device import Buffer
+
+try:
+  from mlx_tinygrad_interop.lib.lease_pool import (
+    MlxToTinygradCopyLeasePool,
+    MlxToTinygradCopyLeasePools,
+    MlxToTinygradLeasePool,
+    MlxToTinygradLeasePools,
+  )
+except ModuleNotFoundError:
+  from lib.lease_pool import MlxToTinygradCopyLeasePool, MlxToTinygradCopyLeasePools, MlxToTinygradLeasePool, MlxToTinygradLeasePools
+
+blackhole: Any = None
+
+DTYPES: dict[str, tuple[Any, Any, np.dtype[Any]]] = {
+  "float16": (mx.float16, dtypes.float16, np.dtype(np.float16)),
+  "float32": (mx.float32, dtypes.float32, np.dtype(np.float32)),
+  "int32": (mx.int32, dtypes.int32, np.dtype(np.int32)),
+  "uint8": (mx.uint8, dtypes.uint8, np.dtype(np.uint8)),
+}
+
+
+class Alternator:
+  def __init__(self, *items: Any):
+    assert items, "Alternator needs at least one item"
+    self.items = items
+    self.index = 0
+
+  def next(self) -> Any:
+    item = self.items[self.index]
+    self.index = (self.index + 1) % len(self.items)
+    return item
+
+
+class BorrowerRing:
+  def __init__(self, *borrowers: Any):
+    assert borrowers, "BorrowerRing needs at least one borrower"
+    self.borrowers = borrowers
+    self.index = 0
+
+  def next(self) -> Any:
+    borrower = self.borrowers[self.index]
+    self.index = (self.index + 1) % len(self.borrowers)
+    return borrower
+
+
+def parse_args() -> argparse.Namespace:
+  parser = argparse.ArgumentParser(description="Benchmark raw tinygrad <-> MLX tensor conversion overhead.")
+  parser.add_argument("--dtype", choices=sorted(DTYPES), default="float32")
+  parser.add_argument("--sizes", default="256,512,1024,2048,4096,7168,8192,16384,32768,65536,262144,1048576",
+                      help="Comma-separated tensor sizes in bytes.")
+  parser.add_argument("--warmup", type=int, default=128)
+  parser.add_argument("--samples", type=int, default=12)
+  parser.add_argument("--min-batch-us", type=float, default=2000.0,
+                      help="Minimum target batch duration per sample in microseconds.")
+  return parser.parse_args()
+
+
+def bytes_view(mv: memoryview) -> memoryview:
+  return mv if mv.format == "B" and mv.ndim == 1 else mv.cast("B")
+
+
+def mlx_dtype_name(dtype: Any) -> str:
+  return repr(dtype).removeprefix("mlx.core.")
+
+
+def tinygrad_zero_copy_memoryview(t: Tensor) -> memoryview:
+  assert t.device == "METAL", f"expected METAL tensor, got {t.device}"
+  buf = cast(Buffer, t.uop.buffer).ensure_allocated()
+  assert t.dtype.base.fmt is not None, f"no buffer format for dtype {t.dtype.base}"
+  return buf.as_memoryview(force_zero_copy=True).cast(t.dtype.base.fmt, t.shape)
+
+
+def tinygrad_from_mlx_legacy(x: Any, tg_dtype: Any) -> Tensor:
+  storage = mx.metal._unsafe_export_storage(x)
+  return Tensor._unsafe_from_metal_buffer(
+    int(storage["mtl_buffer_ptr"]),
+    tuple(storage["shape"]),
+    dtype=tg_dtype,
+    byte_offset=int(storage["offset_bytes"]),
+    buffer_nbytes=int(storage["buffer_nbytes"]),
+    owner=x,
+  )
+
+
+def tinygrad_from_mlx_fast(x: Any, tg_dtype: Any) -> Tensor:
+  storage = mx.metal._unsafe_export_storage(x)
+  return Tensor._unsafe_from_metal_buffer_fast(
+    int(storage["mtl_buffer_ptr"]),
+    tuple(storage["shape"]),
+    dtype=tg_dtype,
+    byte_offset=int(storage["offset_bytes"]),
+    buffer_nbytes=int(storage["buffer_nbytes"]),
+    owner=x,
+  )
+
+
+def tinygrad_from_mlx_single_entry(x: Any, tg_dtype: Any) -> Tensor:
+  return mx.metal._unsafe_to_tinygrad_fast(x, tg_dtype, owner=x)
+
+
+def tinygrad_from_mlx_reuse(x: Any, borrower: Any) -> Tensor:
+  return mx.metal._unsafe_rebind_tinygrad(x, borrower, owner=x)
+
+
+def tinygrad_from_mlx_lease_acquire_release(x: Any, pool: MlxToTinygradLeasePool) -> int:
+  lease = pool.acquire_from_mlx(x)
+  generation = lease.generation
+  lease.release(synchronize=False)
+  return generation
+
+
+def tinygrad_from_mlx_lease_then_use_sum(x: Any, pool: MlxToTinygradLeasePool) -> Tensor:
+  lease = pool.acquire_from_mlx(x)
+  try:
+    return tinygrad_consume_sum(lease.tensor)
+  finally:
+    lease.release(synchronize=False)
+
+
+def tinygrad_from_mlx_scoped_handoff_noop(x: Any, pools: MlxToTinygradLeasePools, tg_dtype: Any) -> int:
+  return pools.run_with_mlx_tensor(x, tg_dtype=tg_dtype, fn=lambda _: 0)
+
+
+def tinygrad_from_mlx_scoped_handoff_then_use_sum(x: Any, pools: MlxToTinygradLeasePools, tg_dtype: Any) -> Tensor:
+  return pools.run_with_mlx_tensor(x, tg_dtype=tg_dtype, fn=lambda t: tinygrad_consume_sum(t))
+
+
+def tinygrad_from_mlx_copy_pool_acquire_release(x: Any, pool: MlxToTinygradCopyLeasePool) -> int:
+  lease = pool.acquire_from_mlx(x)
+  generation = lease.generation
+  lease.release(synchronize=False)
+  return generation
+
+
+def tinygrad_from_mlx_copy_pool_then_use_sum(x: Any, pool: MlxToTinygradCopyLeasePool) -> Tensor:
+  lease = pool.acquire_from_mlx(x)
+  try:
+    return tinygrad_consume_sum(lease.tensor)
+  finally:
+    lease.release(synchronize=False)
+
+
+def tinygrad_from_mlx_scoped_copy_handoff_noop(x: Any, pools: MlxToTinygradCopyLeasePools, tg_dtype: Any) -> int:
+  return pools.run_with_mlx_tensor(x, tg_dtype=tg_dtype, fn=lambda _: 0)
+
+
+def tinygrad_from_mlx_scoped_copy_handoff_then_use_sum(x: Any, pools: MlxToTinygradCopyLeasePools, tg_dtype: Any) -> Tensor:
+  return pools.run_with_mlx_tensor(x, tg_dtype=tg_dtype, fn=lambda t: tinygrad_consume_sum(t))
+
+
+def tinygrad_consume_sum(t: Tensor) -> Tensor:
+  out = (t + 1).sum()
+  out.realize()
+  Device["METAL"].synchronize()
+  return out
+
+
+def mlx_from_tinygrad_maybe_copy(t: Tensor, mx_dtype: Any) -> Any:
+  storage = t._unsafe_metal_storage()
+  return mx.metal._unsafe_array_from_ptr(
+    int(storage["raw_ptr"]),
+    tuple(storage["shape"]),
+    mx_dtype,
+    owner=t,
+  )
+
+
+def mlx_from_tinygrad_alias_only(t: Tensor, mx_dtype: Any) -> Any:
+  storage = t._unsafe_metal_storage()
+  return mx.metal._unsafe_array_from_ptr_alias_only(
+    int(storage["raw_ptr"]),
+    tuple(storage["shape"]),
+    mx_dtype,
+    owner=t,
+  )
+
+
+def tinygrad_from_mlx_copy(x: Any, tg_dtype: Any) -> Tensor:
+  out = Tensor.empty(*tuple(int(dim) for dim in x.shape), device="METAL", dtype=tg_dtype)
+  cast(Buffer, out.uop.buffer).ensure_allocated().copyin(bytes_view(memoryview(x)))
+  return out
+
+
+def mlx_from_tinygrad_copy(t: Tensor) -> Any:
+  return mx.array(tinygrad_zero_copy_memoryview(t))
+
+
+def tinygrad_from_mlx_numpy(x: Any) -> Tensor:
+  out = Tensor(np.array(x, copy=True), device="METAL")
+  out.realize()
+  Device["METAL"].synchronize()
+  return out
+
+
+def mlx_from_tinygrad_numpy(t: Tensor) -> Any:
+  return mx.array(t.numpy())
+
+
+def mlx_export_storage(x: Any) -> Any:
+  return mx.metal._unsafe_export_storage(x)
+
+
+def tinygrad_import_from_storage(storage: Any, tg_dtype: Any, owner: Any) -> Tensor:
+  return Tensor._unsafe_from_metal_buffer(
+    int(storage["mtl_buffer_ptr"]),
+    tuple(storage["shape"]),
+    dtype=tg_dtype,
+    byte_offset=int(storage["offset_bytes"]),
+    buffer_nbytes=int(storage["buffer_nbytes"]),
+    owner=owner,
+  )
+
+
+def tinygrad_import_from_storage_fast(storage: Any, tg_dtype: Any, owner: Any) -> Tensor:
+  return Tensor._unsafe_from_metal_buffer_fast(
+    int(storage["mtl_buffer_ptr"]),
+    tuple(storage["shape"]),
+    dtype=tg_dtype,
+    byte_offset=int(storage["offset_bytes"]),
+    buffer_nbytes=int(storage["buffer_nbytes"]),
+    owner=owner,
+  )
+
+
+def tinygrad_import_from_storage_reuse(storage: Any, owner: Any, borrower: Any) -> Tensor:
+  return borrower.rebind(
+    int(storage["mtl_buffer_ptr"]),
+    owner=owner,
+    shape=tuple(storage["shape"]),
+    dtype_name=mlx_dtype_name(storage["dtype"]),
+    byte_offset=int(storage["offset_bytes"]),
+    buffer_nbytes=int(storage["buffer_nbytes"]),
+  )
+
+
+def tinygrad_export_storage(t: Tensor) -> Any:
+  return t._unsafe_metal_storage()
+
+
+def mlx_import_from_storage(storage: Any, mx_dtype: Any, owner: Any) -> Any:
+  return mx.metal._unsafe_array_from_ptr(
+    int(storage["raw_ptr"]),
+    tuple(storage["shape"]),
+    mx_dtype,
+    owner=owner,
+  )
+
+
+def mlx_import_from_storage_alias_only(storage: Any, mx_dtype: Any, owner: Any) -> Any:
+  return mx.metal._unsafe_array_from_ptr_alias_only(
+    int(storage["raw_ptr"]),
+    tuple(storage["shape"]),
+    mx_dtype,
+    owner=owner,
+  )
+
+
+def bench_callable(fn: Callable[[], Any], warmup: int, samples: int, min_batch_us: float) -> dict[str, float]:
+  global blackhole
+  for _ in range(warmup):
+    blackhole = fn()
+
+  min_batch_ns = int(min_batch_us * 1000.0)
+  iters = 1
+  while True:
+    start = time.perf_counter_ns()
+    for _ in range(iters):
+      blackhole = fn()
+    elapsed = time.perf_counter_ns() - start
+    if elapsed >= min_batch_ns or iters >= (1 << 20):
+      break
+    iters *= 2
+
+  vals_us: list[float] = []
+  for _ in range(samples):
+    start = time.perf_counter_ns()
+    for _ in range(iters):
+      blackhole = fn()
+    elapsed = time.perf_counter_ns() - start
+    vals_us.append(elapsed / iters / 1000.0)
+
+  return {
+    "iters": float(iters),
+    "avg_us": statistics.mean(vals_us),
+    "stddev_us": statistics.stdev(vals_us) if len(vals_us) > 1 else 0.0,
+  }
+
+
+def print_header(dtype_name: str) -> None:
+  print(f"# python={platform.python_version()} platform={platform.platform()}")
+  print(f"# dtype={dtype_name} mlx_metal_available={mx.metal.is_available()} tinygrad_device=METAL")
+  print("# sizes are source tensor sizes in bytes")
+  print("# timed loop excludes source tensor construction and explicit pre-sync, but still includes per-call helper, binding, and wrapper overhead")
+  print("# reported latency is average per-call time with sample standard deviation")
+  print("# rebindable_slot_* rows rebind and return the same tinygrad Tensor object each time; ring rows rotate through multiple such slots")
+  print("# copy_pool_* rows reuse tinygrad-owned destination tensors and copy source bytes into them before release")
+  print("size_bytes,method,direction,avg_us,stddev_us,iters")
+
+
+def main() -> None:
+  args = parse_args()
+  mx_dtype, tg_dtype, np_dtype = DTYPES[args.dtype]
+  sizes = [int(x.strip()) for x in args.sizes.split(",") if x.strip()]
+
+  required = [
+    ("mx.metal._unsafe_export_storage", getattr(mx.metal, "_unsafe_export_storage", None)),
+    ("mx.metal._unsafe_to_tinygrad_fast", getattr(mx.metal, "_unsafe_to_tinygrad_fast", None)),
+    ("mx.metal._unsafe_rebind_tinygrad", getattr(mx.metal, "_unsafe_rebind_tinygrad", None)),
+    ("mx.metal._unsafe_array_from_ptr", getattr(mx.metal, "_unsafe_array_from_ptr", None)),
+    ("mx.metal._unsafe_array_from_ptr_alias_only", getattr(mx.metal, "_unsafe_array_from_ptr_alias_only", None)),
+    ("Tensor._unsafe_from_metal_buffer", getattr(Tensor, "_unsafe_from_metal_buffer", None)),
+    ("Tensor._unsafe_from_metal_buffer_fast", getattr(Tensor, "_unsafe_from_metal_buffer_fast", None)),
+    ("Tensor._unsafe_metal_borrower", getattr(Tensor, "_unsafe_metal_borrower", None)),
+    ("Tensor._unsafe_metal_storage", getattr(Tensor, "_unsafe_metal_storage", None)),
+  ]
+  missing = [name for name, value in required if value is None]
+  if missing:
+    raise RuntimeError(f"Missing required helper(s): {', '.join(missing)}")
+
+  gc.disable()
+  try:
+    print_header(args.dtype)
+    for size_bytes in sizes:
+      if size_bytes <= 0:
+        continue
+      if size_bytes % np_dtype.itemsize != 0:
+        print(f"# skipping size {size_bytes}: not divisible by dtype itemsize {np_dtype.itemsize}", file=sys.stderr)
+        continue
+
+      numel = size_bytes // np_dtype.itemsize
+
+      # Build one realized source tensor on each side. Source creation and the
+      # explicit pre-sync stay outside the timed loop, but per-call helper,
+      # binding, owner-pinning, and wrapper construction still remain inside it.
+      src_mx_pool = [mx.array(np.arange(numel, dtype=np_dtype) + np.array(i, dtype=np_dtype), dtype=mx_dtype) for i in range(4)]
+      src_mx = src_mx_pool[0]
+
+      src_tg = Tensor(np.arange(numel, dtype=np_dtype), device="METAL").realize()
+      Device["METAL"].synchronize()
+
+      mx_storage_pool = [(mlx_export_storage(src), src) for src in src_mx_pool]
+      mx_storage = mx_storage_pool[0][0]
+      tg_storage = tinygrad_export_storage(src_tg)
+      mx_slot_borrower = Tensor._unsafe_metal_borrower(
+        int(mx_storage["mtl_buffer_ptr"]),
+        tuple(mx_storage["shape"]),
+        dtype=tg_dtype,
+        byte_offset=int(mx_storage["offset_bytes"]),
+        buffer_nbytes=int(mx_storage["buffer_nbytes"]),
+        owner=src_mx,
+      )
+      mx_slot_sources = Alternator(*src_mx_pool)
+      mx_slot_storage = Alternator(*mx_storage_pool)
+      mx_ring = BorrowerRing(*[
+        Tensor._unsafe_metal_borrower(
+          int(mx_storage["mtl_buffer_ptr"]),
+          tuple(mx_storage["shape"]),
+          dtype=tg_dtype,
+          byte_offset=int(mx_storage["offset_bytes"]),
+          buffer_nbytes=int(mx_storage["buffer_nbytes"]),
+          owner=src_mx,
+        ) for _ in range(4)
+      ])
+      lease_pool = MlxToTinygradLeasePool.from_mlx(src_mx, tg_dtype=tg_dtype, capacity=4, synchronize_on_release=True)
+      scoped_handoff_pools = MlxToTinygradLeasePools(capacity_per_key=4, synchronize_on_release=True)
+      copy_lease_pool = MlxToTinygradCopyLeasePool.from_mlx(src_mx, tg_dtype=tg_dtype, capacity=4, synchronize_on_release=True)
+      scoped_copy_handoff_pools = MlxToTinygradCopyLeasePools(capacity_per_key=4, synchronize_on_release=True)
+
+      benches: list[tuple[str, str, Callable[[], Any]]] = [
+        ("unsafe_helper_bridge", "mlx_to_tinygrad", lambda s=src_mx: tinygrad_from_mlx_fast(s, tg_dtype)),
+        ("single_entry_bridge", "mlx_to_tinygrad", lambda s=src_mx: tinygrad_from_mlx_single_entry(s, tg_dtype)),
+        ("fresh_wrapper_then_use_sum", "mlx_to_tinygrad", lambda s=src_mx: tinygrad_consume_sum(tinygrad_from_mlx_fast(s, tg_dtype))),
+        ("rebindable_slot_bridge", "mlx_to_tinygrad", lambda alt=mx_slot_sources, b=mx_slot_borrower: tinygrad_from_mlx_reuse(alt.next(), b)),
+        ("rebindable_slot_then_use_sum", "mlx_to_tinygrad",
+         lambda alt=mx_slot_sources, b=mx_slot_borrower: tinygrad_consume_sum(tinygrad_from_mlx_reuse(alt.next(), b))),
+        ("borrower_ring4_bridge", "mlx_to_tinygrad",
+         lambda alt=mx_slot_sources, ring=mx_ring: tinygrad_from_mlx_reuse(alt.next(), ring.next())),
+        ("borrower_ring4_then_use_sum", "mlx_to_tinygrad",
+         lambda alt=mx_slot_sources, ring=mx_ring: tinygrad_consume_sum(tinygrad_from_mlx_reuse(alt.next(), ring.next()))),
+        ("lease_pool_acquire_release", "mlx_to_tinygrad",
+         lambda alt=mx_slot_sources, pool=lease_pool: tinygrad_from_mlx_lease_acquire_release(alt.next(), pool)),
+        ("lease_pool_then_use_sum", "mlx_to_tinygrad",
+         lambda alt=mx_slot_sources, pool=lease_pool: tinygrad_from_mlx_lease_then_use_sum(alt.next(), pool)),
+        ("scoped_handoff_noop", "mlx_to_tinygrad",
+         lambda alt=mx_slot_sources, pools=scoped_handoff_pools, dtype=tg_dtype: tinygrad_from_mlx_scoped_handoff_noop(alt.next(), pools, dtype)),
+        ("scoped_handoff_then_use_sum", "mlx_to_tinygrad",
+         lambda alt=mx_slot_sources, pools=scoped_handoff_pools, dtype=tg_dtype: tinygrad_from_mlx_scoped_handoff_then_use_sum(alt.next(), pools, dtype)),
+        ("copy_pool_acquire_release", "mlx_to_tinygrad",
+         lambda alt=mx_slot_sources, pool=copy_lease_pool: tinygrad_from_mlx_copy_pool_acquire_release(alt.next(), pool)),
+        ("copy_pool_then_use_sum", "mlx_to_tinygrad",
+         lambda alt=mx_slot_sources, pool=copy_lease_pool: tinygrad_from_mlx_copy_pool_then_use_sum(alt.next(), pool)),
+        ("scoped_copy_handoff_noop", "mlx_to_tinygrad",
+         lambda alt=mx_slot_sources, pools=scoped_copy_handoff_pools, dtype=tg_dtype: tinygrad_from_mlx_scoped_copy_handoff_noop(alt.next(), pools, dtype)),
+        ("scoped_copy_handoff_then_use_sum", "mlx_to_tinygrad",
+         lambda alt=mx_slot_sources, pools=scoped_copy_handoff_pools, dtype=tg_dtype: tinygrad_from_mlx_scoped_copy_handoff_then_use_sum(alt.next(), pools, dtype)),
+        ("unsafe_helper_legacy", "mlx_to_tinygrad", lambda s=src_mx: tinygrad_from_mlx_legacy(s, tg_dtype)),
+        ("memoryview_copy", "mlx_to_tinygrad", lambda s=src_mx: tinygrad_from_mlx_copy(s, tg_dtype)),
+        ("numpy_baseline", "mlx_to_tinygrad", lambda s=src_mx: tinygrad_from_mlx_numpy(s)),
+        ("unsafe_helper_bridge", "tinygrad_to_mlx", lambda s=src_tg: mlx_from_tinygrad_alias_only(s, mx_dtype)),
+        ("unsafe_helper_maybe_copy", "tinygrad_to_mlx", lambda s=src_tg: mlx_from_tinygrad_maybe_copy(s, mx_dtype)),
+        ("memoryview_copy", "tinygrad_to_mlx", lambda s=src_tg: mlx_from_tinygrad_copy(s)),
+        ("numpy_baseline", "tinygrad_to_mlx", lambda s=src_tg: mlx_from_tinygrad_numpy(s)),
+        ("export_helper_only", "mlx_to_tinygrad", lambda s=src_mx: mlx_export_storage(s)),
+        ("import_helper_fast_only", "mlx_to_tinygrad", lambda st=mx_storage, s=src_mx: tinygrad_import_from_storage_fast(st, tg_dtype, s)),
+        ("rebindable_slot_import_only", "mlx_to_tinygrad",
+         lambda alt=mx_slot_storage, b=mx_slot_borrower: tinygrad_import_from_storage_reuse(*alt.next(), borrower=b)),
+        ("borrower_ring4_import_only", "mlx_to_tinygrad",
+         lambda alt=mx_slot_storage, ring=mx_ring: tinygrad_import_from_storage_reuse(*alt.next(), borrower=ring.next())),
+        ("import_helper_legacy_only", "mlx_to_tinygrad", lambda st=mx_storage, s=src_mx: tinygrad_import_from_storage(st, tg_dtype, s)),
+        ("export_helper_only", "tinygrad_to_mlx", lambda s=src_tg: tinygrad_export_storage(s)),
+        ("import_helper_only", "tinygrad_to_mlx", lambda st=tg_storage, s=src_tg: mlx_import_from_storage_alias_only(st, mx_dtype, s)),
+        ("import_helper_maybe_copy_only", "tinygrad_to_mlx", lambda st=tg_storage, s=src_tg: mlx_import_from_storage(st, mx_dtype, s)),
+      ]
+
+      for method, direction, fn in benches:
+        try:
+          blackhole = fn()
+        except Exception as exc:
+          print(f"# skipping {size_bytes},{method},{direction}: {exc}", file=sys.stderr)
+          continue
+        stats = bench_callable(fn, warmup=args.warmup, samples=args.samples, min_batch_us=args.min_batch_us)
+        print(
+          f"{size_bytes},{method},{direction},"
+          f"{stats['avg_us']:.3f},{stats['stddev_us']:.3f},{int(stats['iters'])}"
+        )
+  finally:
+    gc.enable()
+
+
+if __name__ == "__main__":
+  main()

mlx_tinygrad_interop/bench_torch_route.py

@@ -0,0 +1,341 @@
+"""Benchmark tinygrad <-> PyTorch <-> MLX bridge routes.
+
+This file intentionally uses pre-existing interop surfaces instead of adding
+new framework patches:
+
+- MLX <-> PyTorch bridge shape based on:
+  https://github.com/vllm-project/vllm-metal/blob/main/vllm_metal/pytorch_backend/tensor_bridge.py
+- PyTorch -> tinygrad via tinygrad's documented Tensor.from_blob runtime interop
+
+The timed loop excludes source tensor construction and explicit pre-sync, but
+still includes the per-call helper and wrapper overhead of the route.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import platform
+import statistics
+import sys
+import time
+from typing import Any, Callable, Literal, cast
+
+import mlx.core as mx
+import numpy as np
+import torch
+from tinygrad import Device, Tensor, dtypes
+from tinygrad.device import Buffer
+from tinygrad.dtype import DType, _from_torch_dtype, _to_torch_dtype
+
+blackhole: Any = None
+
+_MPS_SAFE_SIZE_BYTES = 1 << 30
+
+DTYPES: dict[str, tuple[Any, Any, np.dtype[Any]]] = {
+  "float16": (mx.float16, dtypes.float16, np.dtype(np.float16)),
+  "float32": (mx.float32, dtypes.float32, np.dtype(np.float32)),
+  "int32": (mx.int32, dtypes.int32, np.dtype(np.int32)),
+  "uint8": (mx.uint8, dtypes.uint8, np.dtype(np.uint8)),
+}
+
+MLX_TO_TORCH_DTYPE: dict[mx.Dtype, torch.dtype] = {
+  mx.float32: torch.float32,
+  mx.float16: torch.float16,
+  mx.bfloat16: torch.bfloat16,
+  mx.int32: torch.int32,
+  mx.int64: torch.int64,
+  mx.int16: torch.int16,
+  mx.int8: torch.int8,
+  mx.uint8: torch.uint8,
+  mx.bool_: torch.bool,
+}
+
+TORCH_TO_MLX_DTYPE: dict[torch.dtype, mx.Dtype] = {v: k for k, v in MLX_TO_TORCH_DTYPE.items()}
+
+
+def parse_args() -> argparse.Namespace:
+  parser = argparse.ArgumentParser(description="Benchmark tinygrad <-> PyTorch <-> MLX bridge overhead.")
+  parser.add_argument("--dtype", choices=sorted(DTYPES), default="float32")
+  parser.add_argument("--sizes", default="256,512,1024,2048,4096,7168,8192,16384,32768,65536",
+                      help="Comma-separated tensor sizes in bytes.")
+  parser.add_argument("--warmup", type=int, default=64)
+  parser.add_argument("--samples", type=int, default=8)
+  parser.add_argument("--min-batch-us", type=float, default=2000.0)
+  parser.add_argument("--torch-device", choices=("cpu", "mps", "auto"), default="cpu",
+                      help="Intermediate torch device to use for the route.")
+  return parser.parse_args()
+
+
+def bytes_view(mv: memoryview) -> memoryview:
+  return mv if mv.format == "B" and mv.ndim == 1 else mv.cast("B")
+
+
+def get_torch_device(kind: Literal["cpu", "mps", "auto"] = "auto") -> torch.device:
+  if kind == "cpu":
+    return torch.device("cpu")
+  if kind == "mps":
+    if not torch.backends.mps.is_available():
+      raise RuntimeError("torch MPS backend is not available")
+    return torch.device("mps")
+  if torch.backends.mps.is_available():
+    return torch.device("mps")
+  return torch.device("cpu")
+
+
+def _get_tensor_size_bytes(shape: tuple[int, ...], dtype_itemsize: int) -> int:
+  size = dtype_itemsize
+  for dim in shape:
+    size *= dim
+  return size
+
+
+def sync_mlx() -> None:
+  try:
+    mx.synchronize()
+  except (AttributeError, TypeError):
+    mx.eval(mx.array(0, dtype=mx.int32))
+
+
+def sync_tinygrad() -> None:
+  Device["METAL"].synchronize()
+
+
+def sync_torch(device: torch.device) -> None:
+  if device.type == "mps":
+    torch.mps.synchronize()
+  elif device.type == "cuda":
+    torch.cuda.synchronize()
+
+
+def tinygrad_zero_copy_memoryview(t: Tensor) -> memoryview:
+  assert t.device == "METAL", f"expected METAL tensor, got {t.device}"
+  buf = cast(Buffer, t.uop.buffer).ensure_allocated()
+  return bytes_view(buf.as_memoryview(force_zero_copy=True))
+
+
+def tinygrad_to_torch(tensor: Tensor, *, device: torch.device | Literal["cpu", "mps"] | None = None,
+                      already_contiguous: bool = False) -> torch.Tensor:
+  if device is None:
+    device = get_torch_device("auto")
+  elif isinstance(device, str):
+    device = torch.device(device)
+
+  if not already_contiguous:
+    tensor = tensor.contiguous()
+  tensor = tensor.cast(tensor.dtype.base).realize()
+
+  torch_dtype = _to_torch_dtype(tensor.dtype.base)
+  if torch_dtype is None:
+    raise ValueError(f"Unsupported tinygrad dtype: {tensor.dtype}")
+
+  if tensor.device == "METAL":
+    sync_tinygrad()
+    buffer = tinygrad_zero_copy_memoryview(tensor)
+  else:
+    if tensor.device != "CPU":
+      tensor = tensor.to("CPU").realize()
+    buffer = bytes_view(tensor.data())
+
+  out = torch.frombuffer(buffer, dtype=torch_dtype).reshape(tuple(int(dim) for dim in tensor.shape))
+
+  if device.type == "mps":
+    if _get_tensor_size_bytes(tuple(int(dim) for dim in tensor.shape), tensor.dtype.itemsize) < _MPS_SAFE_SIZE_BYTES:
+      out = out.to(device)
+  elif device.type != "cpu":
+    out = out.to(device)
+
+  return out
+
+
+def torch_to_mlx(tensor: torch.Tensor) -> mx.array:
+  if tensor.device.type != "cpu":
+    sync_torch(tensor.device)
+    tensor = tensor.cpu()
+  tensor = tensor.detach()
+  if tensor.dtype == torch.bfloat16:
+    return mx.array(tensor)
+  return mx.array(tensor.numpy())
+
+
+def mlx_to_torch(array: mx.array, *, device: torch.device | Literal["cpu", "mps"] | None = None,
+                 already_contiguous: bool = False) -> torch.Tensor:
+  if device is None:
+    device = get_torch_device("auto")
+  elif isinstance(device, str):
+    device = torch.device(device)
+
+  torch_dtype = MLX_TO_TORCH_DTYPE.get(array.dtype)
+  if torch_dtype is None:
+    raise ValueError(f"Unsupported MLX dtype: {array.dtype}")
+
+  if not already_contiguous:
+    array = mx.contiguous(array)
+  mx.eval(array)
+  out = torch.frombuffer(memoryview(array), dtype=torch_dtype).reshape(tuple(int(dim) for dim in array.shape))
+
+  if device.type == "mps":
+    if _get_tensor_size_bytes(tuple(int(dim) for dim in array.shape), int(array.dtype.size)) < _MPS_SAFE_SIZE_BYTES:
+      out = out.to(device)
+  elif device.type != "cpu":
+    out = out.to(device)
+
+  return out
+
+
+def torch_to_tinygrad(tensor: torch.Tensor) -> Tensor:
+  tensor = tensor.detach()
+  if not tensor.is_contiguous():
+    tensor = tensor.contiguous()
+
+  if tensor.device.type == "mps":
+    sync_torch(tensor.device)
+    target_device = "METAL"
+  elif tensor.device.type == "cuda":
+    sync_torch(tensor.device)
+    target_device = "CUDA"
+  elif tensor.device.type == "cpu":
+    target_device = "CPU"
+  else:
+    raise ValueError(f"Unsupported torch device: {tensor.device}")
+
+  out = Tensor.from_blob(
+    tensor.data_ptr(),
+    tuple(int(dim) for dim in tensor.shape),
+    dtype=_from_torch_dtype(tensor.dtype),
+    device=target_device,
+  )
+  if out.uop.has_buffer_identity():
+    setattr(cast(Buffer, out.uop.buffer).base, "_external_owner", tensor)
+  return out
+
+
+def tinygrad_to_mlx_via_torch(tensor: Tensor, *, torch_device: torch.device) -> mx.array:
+  return torch_to_mlx(tinygrad_to_torch(tensor, device=torch_device))
+
+
+def mlx_to_tinygrad_via_torch(array: mx.array, *, torch_device: torch.device) -> Tensor:
+  return torch_to_tinygrad(mlx_to_torch(array, device=torch_device))
+
+
+def tinygrad_to_mlx_direct(tensor: Tensor) -> mx.array:
+  return mx.array(tinygrad_zero_copy_memoryview(tensor))
+
+
+def mlx_to_tinygrad_direct(array: mx.array, tg_dtype: DType) -> Tensor:
+  storage = mx.metal._unsafe_export_storage(array)
+  return Tensor._unsafe_from_metal_buffer_fast(
+    int(storage["mtl_buffer_ptr"]),
+    tuple(storage["shape"]),
+    dtype=tg_dtype,
+    byte_offset=int(storage["offset_bytes"]),
+    buffer_nbytes=int(storage["buffer_nbytes"]),
+    owner=array,
+  )
+
+
+def bench_callable(fn: Callable[[], Any], warmup: int, samples: int, min_batch_us: float) -> dict[str, float]:
+  global blackhole
+  for _ in range(warmup):
+    blackhole = fn()
+
+  min_batch_ns = int(min_batch_us * 1000.0)
+  iters = 1
+  while True:
+    start = time.perf_counter_ns()
+    for _ in range(iters):
+      blackhole = fn()
+    elapsed = time.perf_counter_ns() - start
+    if elapsed >= min_batch_ns or iters >= (1 << 20):
+      break
+    iters *= 2
+
+  vals_us: list[float] = []
+  for _ in range(samples):
+    start = time.perf_counter_ns()
+    for _ in range(iters):
+      blackhole = fn()
+    elapsed = time.perf_counter_ns() - start
+    vals_us.append(elapsed / iters / 1000.0)
+
+  return {
+    "iters": float(iters),
+    "avg_us": statistics.mean(vals_us),
+    "stddev_us": statistics.stdev(vals_us) if len(vals_us) > 1 else 0.0,
+  }
+
+
+def assert_equal(name: str, actual: np.ndarray, expected: np.ndarray) -> None:
+  if np.issubdtype(expected.dtype, np.floating):
+    np.testing.assert_allclose(actual, expected, rtol=5e-5 if expected.dtype == np.float32 else 5e-3, atol=1e-5 if expected.dtype == np.float32 else 5e-3, err_msg=name)
+  else:
+    np.testing.assert_array_equal(actual, expected, err_msg=name)
+
+
+def print_header(dtype_name: str, torch_device: torch.device) -> None:
+  print(f"# python={platform.python_version()} platform={platform.platform()}")
+  print(f"# dtype={dtype_name} mlx_metal_available={mx.metal.is_available()} tinygrad_device=METAL torch_device={torch_device}")
+  print("# sizes are source tensor sizes in bytes")
+  print("# timed loop excludes source tensor construction and explicit pre-sync, but still includes helper and wrapper overhead")
+  print("# reported latency is average per-call time with sample standard deviation")
+  print("# via_torch_route rows use pre-existing tinygrad<->torch and mlx<->torch bridges without framework patches")
+  print("size_bytes,method,direction,avg_us,stddev_us,iters")
+
+
+def main() -> None:
+  args = parse_args()
+  if not mx.metal.is_available():
+    raise RuntimeError("MLX Metal is not available")
+  torch_device = get_torch_device(cast(Literal["cpu", "mps", "auto"], args.torch_device))
+  mx_dtype, tg_dtype, np_dtype = DTYPES[args.dtype]
+  sizes = [int(x.strip()) for x in args.sizes.split(",") if x.strip()]
+
+  gc.disable()
+  try:
+    print_header(args.dtype, torch_device)
+    for size_bytes in sizes:
+      if size_bytes <= 0:
+        continue
+      if size_bytes % np_dtype.itemsize != 0:
+        print(f"# skipping size {size_bytes}: not divisible by dtype itemsize {np_dtype.itemsize}", file=sys.stderr)
+        continue
+
+      numel = size_bytes // np_dtype.itemsize
+      values = np.arange(numel, dtype=np_dtype)
+      src_tg = Tensor(values, device="METAL", dtype=tg_dtype).realize()
+      src_mx = mx.array(values, dtype=mx_dtype)
+      sync_tinygrad()
+      sync_mlx()
+
+      # correctness checks stay outside the timed loop
+      assert_equal("tinygrad->mlx via torch raw", np.array(tinygrad_to_mlx_via_torch(src_tg, torch_device=torch_device)), values)
+      assert_equal("mlx->tinygrad via torch raw", mlx_to_tinygrad_via_torch(src_mx, torch_device=torch_device).numpy(), values)
+
+      benches: list[tuple[str, str, Callable[[], Any]]] = [
+        ("via_torch_route", "tinygrad_to_mlx", lambda s=src_tg, td=torch_device: tinygrad_to_mlx_via_torch(s, torch_device=td)),
+        ("via_torch_route", "mlx_to_tinygrad", lambda s=src_mx, td=torch_device: mlx_to_tinygrad_via_torch(s, torch_device=td)),
+        ("bridge_half", "tinygrad_to_torch", lambda s=src_tg, td=torch_device: tinygrad_to_torch(s, device=td)),
+        ("bridge_half", "torch_to_mlx", lambda s=src_tg, td=torch_device: torch_to_mlx(tinygrad_to_torch(s, device=td))),
+        ("bridge_half", "mlx_to_torch", lambda s=src_mx, td=torch_device: mlx_to_torch(s, device=td)),
+        ("bridge_half", "torch_to_tinygrad", lambda s=src_mx, td=torch_device: torch_to_tinygrad(mlx_to_torch(s, device=td))),
+        ("direct_baseline", "tinygrad_to_mlx", lambda s=src_tg: tinygrad_to_mlx_direct(s)),
+        ("direct_baseline", "mlx_to_tinygrad", lambda s=src_mx, td=tg_dtype: mlx_to_tinygrad_direct(s, td)),
+      ]
+
+      for method, direction, fn in benches:
+        try:
+          blackhole = fn()
+        except Exception as exc:
+          print(f"# skipping {size_bytes},{method},{direction}: {exc}", file=sys.stderr)
+          continue
+        stats = bench_callable(fn, warmup=args.warmup, samples=args.samples, min_batch_us=args.min_batch_us)
+        print(
+          f"{size_bytes},{method},{direction},"
+          f"{stats['avg_us']:.3f},{stats['stddev_us']:.3f},{int(stats['iters'])}"
+        )
+  finally:
+    gc.enable()
+
+
+if __name__ == "__main__":
+  main()

mlx_tinygrad_interop/lease_pool.py

@@ -0,0 +1,7 @@
+"""Compatibility shim for the moved lease-pool implementation.
+
+Reusable interop code now lives under `mlx_tinygrad_interop.lib`.
+Benchmarks and tests remain at the package top level.
+"""
+
+from mlx_tinygrad_interop.lib.lease_pool import *  # noqa: F401,F403

mlx_tinygrad_interop/lib/__init__.py

@@ -0,0 +1,31 @@
+"""Reusable MLX <-> tinygrad interop helpers.
+
+Benchmarks, stress harnesses, and tests stay in the top-level
+`mlx_tinygrad_interop/` package. Reusable bridge and lease-pool code lives
+under `mlx_tinygrad_interop/lib/`.
+"""
+
+from mlx_tinygrad_interop.lib.lease_pool import (
+  MlxToTinygradCopyKey,
+  MlxToTinygradCopyLeasePool,
+  MlxToTinygradCopyLeasePools,
+  MlxToTinygradLease,
+  MlxToTinygradLeaseKey,
+  MlxToTinygradLeasePool,
+  MlxToTinygradLeasePools,
+)
+from mlx_tinygrad_interop.lib.tensor_bridge import mlx_to_tinygrad, sync_mlx, sync_tinygrad, tinygrad_to_mlx
+
+__all__ = [
+  "MlxToTinygradCopyKey",
+  "MlxToTinygradCopyLeasePool",
+  "MlxToTinygradCopyLeasePools",
+  "MlxToTinygradLease",
+  "MlxToTinygradLeaseKey",
+  "MlxToTinygradLeasePool",
+  "MlxToTinygradLeasePools",
+  "mlx_to_tinygrad",
+  "sync_mlx",
+  "sync_tinygrad",
+  "tinygrad_to_mlx",
+]

mlx_tinygrad_interop/lib/lease_pool.py

@@ -0,0 +1,447 @@
+from __future__ import annotations
+
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Any, cast
+
+import mlx.core as mx
+from tinygrad import Device, Tensor
+from tinygrad.device import Buffer
+from tinygrad.dtype import DTypeLike, to_dtype
+from tinygrad.tensor import all_tensors
+
+
+def _export_realized_storage(array: Any) -> dict[str, Any]:
+  mx.eval(array)
+  return mx.metal._unsafe_export_storage(array)
+
+
+def _mlx_dtype_name(dtype: Any) -> str:
+  return repr(dtype).removeprefix("mlx.core.")
+
+
+def _bytes_view(mv: memoryview) -> memoryview:
+  return mv if mv.format == "B" and mv.ndim == 1 else mv.cast("B")
+
+
+def _iter_tensors(obj: Any, seen: set[int] | None = None):
+  if seen is None: seen = set()
+  obj_id = id(obj)
+  if obj_id in seen: return
+  seen.add(obj_id)
+
+  if isinstance(obj, Tensor):
+    yield obj
+    return
+  if isinstance(obj, dict):
+    for value in obj.values():
+      yield from _iter_tensors(value, seen)
+    return
+  if isinstance(obj, (list, tuple, set, frozenset)):
+    for value in obj:
+      yield from _iter_tensors(value, seen)
+
+
+def _snapshot_live_tensors() -> dict[int, Tensor]:
+  return {id(t): t for tref in list(all_tensors) if (t := tref()) is not None}
+
+
+def _tensor_base_buffer(t: Tensor) -> Buffer | None:
+  if not t.uop.has_buffer_identity(): return None
+  try:
+    return cast(Buffer, t.uop.buffer).base
+  except Exception:
+    return None
+
+
+def _uop_depends_on(uop: Any, target: Any) -> bool:
+  seen: set[Any] = set()
+  stack = [uop]
+  while stack:
+    cur = stack.pop()
+    if cur is target: return True
+    if cur in seen: continue
+    seen.add(cur)
+    stack.extend(cur.src)
+  return False
+
+
+def _reject_scoped_tensor_escapes(borrowed: Tensor, returned_tensors: tuple[Tensor, ...], pre_live_ids: set[int]) -> None:
+  borrowed_base = _tensor_base_buffer(borrowed)
+
+  if any(t is borrowed for t in returned_tensors):
+    raise RuntimeError("callback returned the borrowed tensor directly; return a copied or independently realized result instead")
+  if borrowed_base is not None and any(_tensor_base_buffer(t) is borrowed_base for t in returned_tensors):
+    raise RuntimeError("callback returned tensor(s) that still alias the borrowed slot; copy or realize independent storage before returning")
+
+  for tensor in returned_tensors:
+    tensor.realize()
+
+  returned_ids = {id(t) for t in returned_tensors}
+  escaped: list[Tensor] = []
+  for tensor_id, tensor in _snapshot_live_tensors().items():
+    if tensor_id in pre_live_ids or tensor_id in returned_ids or tensor is borrowed:
+      continue
+    if borrowed_base is not None and _tensor_base_buffer(tensor) is borrowed_base:
+      escaped.append(tensor)
+      continue
+    if _uop_depends_on(tensor.uop, borrowed.uop):
+      escaped.append(tensor)
+  if escaped:
+    raise RuntimeError(
+      "callback leaked tensor(s) derived from the borrowed tensor; only independent realized outputs may escape the callback"
+    )
+
+
+def _run_with_scoped_lease(lease: "MlxToTinygradLease", fn, *, synchronize_on_release: bool | None = None):
+  borrowed = lease.tensor
+  pre_live_ids = set(_snapshot_live_tensors())
+  try:
+    result = fn(borrowed)
+    returned_tensors = tuple(_iter_tensors(result))
+    _reject_scoped_tensor_escapes(borrowed, returned_tensors, pre_live_ids)
+    return result
+  finally:
+    if not lease._released:
+      lease.release(synchronize=synchronize_on_release)
+
+
+def _evict_lru_idle_pool(pools: "OrderedDict[Any, Any]", *, max_pools: int | None) -> None:
+  if max_pools is None or len(pools) < max_pools: return
+  for key, pool in list(pools.items()):
+    if pool.in_flight == 0:
+      del pools[key]
+      return
+  raise RuntimeError(f"pool cache is full ({max_pools} keys) and every pool is still in use")
+
+
+@dataclass(frozen=True)
+class MlxToTinygradLeaseKey:
+  shape: tuple[int, ...]
+  mlx_dtype_name: str
+  tinygrad_dtype_name: str
+  byte_offset: int
+
+  @staticmethod
+  def from_storage(storage: dict[str, Any], *, tg_dtype: DTypeLike) -> "MlxToTinygradLeaseKey":
+    return MlxToTinygradLeaseKey(
+      shape=tuple(int(dim) for dim in storage["shape"]),
+      mlx_dtype_name=_mlx_dtype_name(storage["dtype"]),
+      tinygrad_dtype_name=to_dtype(tg_dtype).base.name,
+      byte_offset=int(storage["offset_bytes"]),
+    )
+
+
+@dataclass(frozen=True)
+class MlxToTinygradCopyKey:
+  shape: tuple[int, ...]
+  mlx_dtype_name: str
+  tinygrad_dtype_name: str
+
+  @staticmethod
+  def from_storage(storage: dict[str, Any], *, tg_dtype: DTypeLike) -> "MlxToTinygradCopyKey":
+    return MlxToTinygradCopyKey(
+      shape=tuple(int(dim) for dim in storage["shape"]),
+      mlx_dtype_name=_mlx_dtype_name(storage["dtype"]),
+      tinygrad_dtype_name=to_dtype(tg_dtype).base.name,
+    )
+
+
+@dataclass
+class _AliasLeaseSlot:
+  borrower: Any
+  tensor: Tensor
+  generation: int = 0
+  in_use: bool = False
+
+
+@dataclass
+class _CopyLeaseSlot:
+  tensor: Tensor
+  generation: int = 0
+  in_use: bool = False
+
+
+class MlxToTinygradLease:
+  __slots__ = ("_pool", "_slot_index", "_generation", "_tensor", "_released")
+
+  def __init__(self, pool: Any, slot_index: int, generation: int, tensor: Tensor):
+    self._pool, self._slot_index, self._generation = pool, slot_index, generation
+    self._tensor, self._released = tensor, False
+
+  @property
+  def generation(self) -> int: return self._generation
+
+  @property
+  def key(self) -> Any: return self._pool.key
+
+  @property
+  def tensor(self) -> Tensor:
+    # This is the unsafe low-level lease surface. The preferred production API
+    # is `run_with_mlx_tensor(...)`, which scopes use and release together.
+    if self._released: raise RuntimeError("lease already released")
+    return self._tensor
+
+  def release(self, *, synchronize: bool | None = None) -> None:
+    if self._released: raise RuntimeError("lease already released")
+    self._pool._release(self._slot_index, self._generation, synchronize=synchronize)
+    self._released = True
+
+  def __enter__(self) -> "MlxToTinygradLease": return self
+
+  def __exit__(self, exc_type, exc, tb) -> bool:
+    if not self._released: self.release()
+    return False
+
+
+class MlxToTinygradLeasePool:
+  __slots__ = ("key", "tg_dtype", "capacity", "_slots", "_next_slot", "_synchronize_on_release")
+
+  def __init__(self, *, key: MlxToTinygradLeaseKey, tg_dtype: DTypeLike, template_storage: dict[str, Any], template_owner: Any,
+               capacity: int = 4, synchronize_on_release: bool = True):
+    if capacity <= 0: raise ValueError(f"capacity must be positive, got {capacity}")
+    self.key, self.tg_dtype, self.capacity = key, to_dtype(tg_dtype), capacity
+    self._next_slot, self._synchronize_on_release = 0, synchronize_on_release
+    self._slots = [
+      _AliasLeaseSlot(
+        borrower := Tensor._unsafe_metal_borrower(
+          int(template_storage["mtl_buffer_ptr"]),
+          tuple(template_storage["shape"]),
+          dtype=self.tg_dtype,
+          byte_offset=int(template_storage["offset_bytes"]),
+          buffer_nbytes=int(template_storage["buffer_nbytes"]),
+          owner=template_owner,
+        ),
+        borrower.tensor,
+      ) for _ in range(capacity)
+    ]
+
+  @classmethod
+  def from_mlx(cls, array: Any, *, tg_dtype: DTypeLike, owner: Any | None = None, capacity: int = 4,
+               synchronize_on_release: bool = True) -> "MlxToTinygradLeasePool":
+    storage = _export_realized_storage(array)
+    owner_obj = array if owner is None else owner
+    return cls(
+      key=MlxToTinygradLeaseKey.from_storage(storage, tg_dtype=tg_dtype),
+      tg_dtype=tg_dtype,
+      template_storage=storage,
+      template_owner=owner_obj,
+      capacity=capacity,
+      synchronize_on_release=synchronize_on_release,
+    )
+
+  def _next_available_slot(self) -> tuple[int, _AliasLeaseSlot]:
+    for _ in range(self.capacity):
+      slot_index = self._next_slot
+      self._next_slot = (self._next_slot + 1) % self.capacity
+      slot = self._slots[slot_index]
+      if not slot.in_use: return slot_index, slot
+    raise RuntimeError(
+      f"all {self.capacity} lease slots for key={self.key} are still in use; "
+      "release leases or increase pool capacity"
+    )
+
+  def _validate_storage(self, storage: dict[str, Any]) -> None:
+    incoming = MlxToTinygradLeaseKey.from_storage(storage, tg_dtype=self.tg_dtype)
+    if incoming != self.key:
+      raise ValueError(f"pool key mismatch: expected {self.key}, got {incoming}")
+
+  def acquire_from_storage(self, storage: dict[str, Any], *, owner: Any) -> MlxToTinygradLease:
+    self._validate_storage(storage)
+    slot_index, slot = self._next_available_slot()
+    tensor = slot.borrower.rebind(
+      int(storage["mtl_buffer_ptr"]),
+      owner=owner,
+      shape=tuple(storage["shape"]),
+      dtype_name=_mlx_dtype_name(storage["dtype"]),
+      byte_offset=int(storage["offset_bytes"]),
+      buffer_nbytes=int(storage["buffer_nbytes"]),
+    )
+    slot.tensor = tensor
+    slot.generation += 1
+    slot.in_use = True
+    return MlxToTinygradLease(self, slot_index, slot.generation, tensor)
+
+  def acquire_from_mlx(self, array: Any, *, owner: Any | None = None) -> MlxToTinygradLease:
+    storage = _export_realized_storage(array)
+    return self.acquire_from_storage(storage, owner=array if owner is None else owner)
+
+  def run_with_mlx_tensor(self, array: Any, fn, *, owner: Any | None = None,
+                          synchronize_on_release: bool | None = None):
+    return _run_with_scoped_lease(
+      self.acquire_from_mlx(array, owner=owner),
+      fn,
+      synchronize_on_release=synchronize_on_release,
+    )
+
+  def _release(self, slot_index: int, generation: int, *, synchronize: bool | None = None) -> None:
+    slot = self._slots[slot_index]
+    if not slot.in_use: raise RuntimeError(f"slot {slot_index} is not currently leased")
+    if slot.generation != generation:
+      raise RuntimeError(f"stale lease generation for slot {slot_index}: expected {slot.generation}, got {generation}")
+    do_synchronize = synchronize if synchronize is not None else self._synchronize_on_release
+    if do_synchronize:
+      Device["METAL"].synchronize()
+      slot.borrower.clear_owner()
+    slot.in_use = False
+
+  @property
+  def in_flight(self) -> int: return sum(int(slot.in_use) for slot in self._slots)
+
+
+class MlxToTinygradLeasePools:
+  __slots__ = ("capacity_per_key", "max_pools", "_synchronize_on_release", "_pools")
+
+  def __init__(self, *, capacity_per_key: int = 4, max_pools: int | None = 64, synchronize_on_release: bool = True):
+    if capacity_per_key <= 0: raise ValueError(f"capacity_per_key must be positive, got {capacity_per_key}")
+    if max_pools is not None and max_pools <= 0: raise ValueError(f"max_pools must be positive, got {max_pools}")
+    self.capacity_per_key, self.max_pools, self._synchronize_on_release = capacity_per_key, max_pools, synchronize_on_release
+    self._pools: OrderedDict[MlxToTinygradLeaseKey, MlxToTinygradLeasePool] = OrderedDict()
+
+  def _get_or_create_pool(self, storage: dict[str, Any], *, tg_dtype: DTypeLike, owner_obj: Any) -> MlxToTinygradLeasePool:
+    key = MlxToTinygradLeaseKey.from_storage(storage, tg_dtype=tg_dtype)
+    pool = self._pools.get(key)
+    if pool is not None:
+      self._pools.move_to_end(key)
+      return pool
+    _evict_lru_idle_pool(self._pools, max_pools=self.max_pools)
+    pool = MlxToTinygradLeasePool(
+      key=key,
+      tg_dtype=tg_dtype,
+      template_storage=storage,
+      template_owner=owner_obj,
+      capacity=self.capacity_per_key,
+      synchronize_on_release=self._synchronize_on_release,
+    )
+    self._pools[key] = pool
+    return pool
+
+  def acquire_from_mlx(self, array: Any, *, tg_dtype: DTypeLike, owner: Any | None = None) -> MlxToTinygradLease:
+    storage = _export_realized_storage(array)
+    owner_obj = array if owner is None else owner
+    pool = self._get_or_create_pool(storage, tg_dtype=tg_dtype, owner_obj=owner_obj)
+    return pool.acquire_from_storage(storage, owner=owner_obj)
+
+  def run_with_mlx_tensor(self, array: Any, *, tg_dtype: DTypeLike, fn, owner: Any | None = None,
+                          synchronize_on_release: bool | None = None):
+    return _run_with_scoped_lease(
+      self.acquire_from_mlx(array, tg_dtype=tg_dtype, owner=owner),
+      fn,
+      synchronize_on_release=synchronize_on_release,
+    )
+
+  @property
+  def pool_count(self) -> int: return len(self._pools)
+
+  def get_pool(self, key: MlxToTinygradLeaseKey) -> MlxToTinygradLeasePool | None:
+    return self._pools.get(key)
+
+
+class MlxToTinygradCopyLeasePool:
+  __slots__ = ("key", "tg_dtype", "capacity", "_slots", "_next_slot", "_synchronize_on_release")
+
+  def __init__(self, *, key: MlxToTinygradCopyKey, tg_dtype: DTypeLike, capacity: int = 4,
+               synchronize_on_release: bool = True):
+    if capacity <= 0: raise ValueError(f"capacity must be positive, got {capacity}")
+    self.key, self.tg_dtype, self.capacity = key, to_dtype(tg_dtype), capacity
+    self._next_slot, self._synchronize_on_release = 0, synchronize_on_release
+    self._slots = [_CopyLeaseSlot(Tensor.empty(*key.shape, dtype=self.tg_dtype, device="METAL").realize()) for _ in range(capacity)]
+
+  @classmethod
+  def from_mlx(cls, array: Any, *, tg_dtype: DTypeLike, capacity: int = 4,
+               synchronize_on_release: bool = True) -> "MlxToTinygradCopyLeasePool":
+    storage = _export_realized_storage(array)
+    return cls(
+      key=MlxToTinygradCopyKey.from_storage(storage, tg_dtype=tg_dtype),
+      tg_dtype=tg_dtype,
+      capacity=capacity,
+      synchronize_on_release=synchronize_on_release,
+    )
+
+  def _next_available_slot(self) -> tuple[int, _CopyLeaseSlot]:
+    for _ in range(self.capacity):
+      slot_index = self._next_slot
+      self._next_slot = (self._next_slot + 1) % self.capacity
+      slot = self._slots[slot_index]
+      if not slot.in_use: return slot_index, slot
+    raise RuntimeError(
+      f"all {self.capacity} copy slots for key={self.key} are still in use; "
+      "release leases or increase pool capacity"
+    )
+
+  def _validate_storage(self, storage: dict[str, Any]) -> None:
+    incoming = MlxToTinygradCopyKey.from_storage(storage, tg_dtype=self.tg_dtype)
+    if incoming != self.key:
+      raise ValueError(f"copy pool key mismatch: expected {self.key}, got {incoming}")
+
+  def acquire_from_mlx(self, array: Any) -> MlxToTinygradLease:
+    storage = _export_realized_storage(array)
+    self._validate_storage(storage)
+    slot_index, slot = self._next_available_slot()
+    cast(Buffer, slot.tensor.uop.buffer).ensure_allocated().copyin(_bytes_view(memoryview(array)))
+    slot.generation += 1
+    slot.in_use = True
+    return MlxToTinygradLease(self, slot_index, slot.generation, slot.tensor)
+
+  def run_with_mlx_tensor(self, array: Any, fn, *, synchronize_on_release: bool | None = None):
+    return _run_with_scoped_lease(
+      self.acquire_from_mlx(array),
+      fn,
+      synchronize_on_release=synchronize_on_release,
+    )
+
+  def _release(self, slot_index: int, generation: int, *, synchronize: bool | None = None) -> None:
+    slot = self._slots[slot_index]
+    if not slot.in_use: raise RuntimeError(f"slot {slot_index} is not currently leased")
+    if slot.generation != generation:
+      raise RuntimeError(f"stale lease generation for slot {slot_index}: expected {slot.generation}, got {generation}")
+    do_synchronize = synchronize if synchronize is not None else self._synchronize_on_release
+    if do_synchronize:
+      Device["METAL"].synchronize()
+    slot.in_use = False
+
+  @property
+  def in_flight(self) -> int: return sum(int(slot.in_use) for slot in self._slots)
+
+
+class MlxToTinygradCopyLeasePools:
+  __slots__ = ("capacity_per_key", "max_pools", "_synchronize_on_release", "_pools")
+
+  def __init__(self, *, capacity_per_key: int = 4, max_pools: int | None = 64, synchronize_on_release: bool = True):
+    if capacity_per_key <= 0: raise ValueError(f"capacity_per_key must be positive, got {capacity_per_key}")
+    if max_pools is not None and max_pools <= 0: raise ValueError(f"max_pools must be positive, got {max_pools}")
+    self.capacity_per_key, self.max_pools, self._synchronize_on_release = capacity_per_key, max_pools, synchronize_on_release
+    self._pools: OrderedDict[MlxToTinygradCopyKey, MlxToTinygradCopyLeasePool] = OrderedDict()
+
+  def _get_or_create_pool(self, storage: dict[str, Any], *, tg_dtype: DTypeLike) -> MlxToTinygradCopyLeasePool:
+    key = MlxToTinygradCopyKey.from_storage(storage, tg_dtype=tg_dtype)
+    pool = self._pools.get(key)
+    if pool is not None:
+      self._pools.move_to_end(key)
+      return pool
+    _evict_lru_idle_pool(self._pools, max_pools=self.max_pools)
+    pool = MlxToTinygradCopyLeasePool(
+      key=key,
+      tg_dtype=tg_dtype,
+      capacity=self.capacity_per_key,
+      synchronize_on_release=self._synchronize_on_release,
+    )
+    self._pools[key] = pool
+    return pool
+
+  def acquire_from_mlx(self, array: Any, *, tg_dtype: DTypeLike) -> MlxToTinygradLease:
+    storage = _export_realized_storage(array)
+    return self._get_or_create_pool(storage, tg_dtype=tg_dtype).acquire_from_mlx(array)
+
+  def run_with_mlx_tensor(self, array: Any, *, tg_dtype: DTypeLike, fn, synchronize_on_release: bool | None = None):
+    return _run_with_scoped_lease(
+      self.acquire_from_mlx(array, tg_dtype=tg_dtype),
+      fn,
+      synchronize_on_release=synchronize_on_release,
+    )
+
+  @property
+  def pool_count(self) -> int: return len(self._pools)
+
+  def get_pool(self, key: MlxToTinygradCopyKey) -> MlxToTinygradCopyLeasePool | None:
+    return self._pools.get(key)

mlx_tinygrad_interop/lib/tensor_bridge.py

@@ -0,0 +1,128 @@
+"""Tensor bridge between tinygrad and MLX.
+
+Baseline implementation modelled after:
+https://github.com/vllm-project/vllm-metal/blob/main/vllm_metal/pytorch_backend/tensor_bridge.py
+
+This module currently implements the `tinygrad -> MLX` direction only.
+`MLX -> tinygrad` is left stubbed on purpose.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Literal, cast
+
+import mlx.core as mx
+from tinygrad.tensor import Tensor
+from tinygrad.device import Buffer, Device
+from tinygrad.dtype import DType, dtypes
+
+logger = logging.getLogger(__name__)
+
+# MPS has a 4GB (2^32 bytes) limit for MPSTemporaryNDArray allocations.
+# Metal may allocate multiple temporary buffers internally, so we use a
+# conservative threshold of 1GB to avoid hitting the limit.
+# See: https://github.com/anthropics/vllm-metal/issues/43
+_MPS_SAFE_SIZE_BYTES = 1 << 30  # 1GB
+
+# MLX to tinygrad dtype mapping
+MLX_TO_TINYGRAD_DTYPE: dict[mx.Dtype, DType] = {
+    mx.float32: dtypes.float32,
+    mx.float16: dtypes.float16,
+    mx.bfloat16: dtypes.bfloat16,
+    mx.int32: dtypes.int32,
+    mx.int64: dtypes.int64,
+    mx.int16: dtypes.int16,
+    mx.int8: dtypes.int8,
+    mx.uint8: dtypes.uint8,
+    mx.bool_: dtypes.bool,
+}
+
+# tinygrad to MLX dtype mapping
+TINYGRAD_TO_MLX_DTYPE: dict[DType, mx.Dtype] = {
+  v: k for k, v in MLX_TO_TINYGRAD_DTYPE.items()
+}
+
+
+def _get_tensor_size_bytes(tensor: Tensor) -> int:
+  """Calculate the size of a tinygrad tensor in bytes."""
+  return tensor.numel() * tensor.dtype.itemsize
+
+
+def _get_buffer_view(tensor: Tensor, *, already_contiguous: bool = False) -> memoryview:
+  """Expose a tinygrad tensor as a Python buffer.
+
+  For METAL tensors on Apple Silicon, this uses tinygrad's zero-copy
+  `as_memoryview(force_zero_copy=True)` path after forcing realization and a
+  Metal synchronize. For other devices, it falls back to the standard
+  `Tensor.data()` path.
+  """
+  tensor = tensor.cast(tensor.dtype.base)
+  if not already_contiguous:
+    tensor = tensor.contiguous()
+
+  if tensor.device == "METAL":
+    tensor = tensor.realize()
+    sync_tinygrad()
+    if tensor.dtype.base.fmt is None:
+      raise ValueError(f"Unsupported tinygrad dtype for memoryview bridge: {tensor.dtype}")
+    buf = cast(Buffer, tensor.uop.buffer).ensure_allocated()
+    return buf.as_memoryview(force_zero_copy=True).cast(tensor.dtype.base.fmt, tensor.shape)
+
+  if tensor.device != "CPU":
+    tensor = tensor.to("CPU").realize()
+
+  return tensor.data()
+
+
+def tinygrad_to_mlx(tensor: Tensor, *, already_contiguous: bool = False) -> mx.array:
+  """Convert a tinygrad tensor to an MLX array.
+
+  Uses a buffer-protocol / memoryview path when possible. In current MLX this
+  still creates a fresh MLX array rather than aliasing the tinygrad buffer, but
+  it is the closest analogue to the reference vLLM bridge's public shape.
+
+  Args:
+    tensor: tinygrad tensor
+    already_contiguous: Skip the contiguity step if the tensor is already known
+      to be dense row-major contiguous.
+
+  Returns:
+    MLX array with the same logical values.
+  """
+  if tensor.dtype.base not in TINYGRAD_TO_MLX_DTYPE:
+    raise ValueError(f"Unsupported tinygrad dtype: {tensor.dtype}")
+
+  buffer = _get_buffer_view(tensor, already_contiguous=already_contiguous)
+  array = mx.array(buffer)
+  if array.dtype != TINYGRAD_TO_MLX_DTYPE[tensor.dtype.base]:
+    array = array.astype(TINYGRAD_TO_MLX_DTYPE[tensor.dtype.base])
+  return array
+
+
+def mlx_to_tinygrad(array: mx.array) -> Tensor:
+  """Convert an MLX array to a tinygrad tensor.
+
+  This direction is intentionally left stubbed here. The current repo already
+  carries more specialized MLX -> tinygrad experiments in the lease-pool and
+  benchmark helpers, and this baseline bridge module is only meant to mirror
+  the public shape of the vLLM bridge for `tinygrad -> MLX`.
+  """
+  raise NotImplementedError("mlx_to_tinygrad() is intentionally stubbed in this baseline bridge")
+
+
+def sync_mlx() -> None:
+  """Synchronize MLX operations."""
+  try:
+    mx.synchronize()
+  except (AttributeError, TypeError):
+    mx.eval(mx.array(0, dtype=mx.int32))
+
+
+def sync_tinygrad() -> None:
+  """Synchronize tinygrad METAL operations."""
+  try:
+    Device["METAL"].synchronize()
+  except Exception:
+    logger.debug("tinygrad METAL synchronize unavailable", exc_info=True)
+

mlx_tinygrad_interop/lib/tensor_bridge_vllm.py

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tensor bridge between MLX and PyTorch.
+
+Provides zero-copy conversion when possible using Apple Silicon's unified memory.
+"""
+
+import logging
+from typing import Literal
+
+import mlx.core as mx
+import torch
+
+logger = logging.getLogger(__name__)
+
+# MPS has a 4GB (2^32 bytes) limit for MPSTemporaryNDArray allocations.
+# Metal may allocate multiple temporary buffers internally, so we use a
+# conservative threshold of 1GB to avoid hitting the limit.
+# See: https://github.com/anthropics/vllm-metal/issues/43
+_MPS_SAFE_SIZE_BYTES = 1 << 30  # 1GB
+
+# MLX to PyTorch dtype mapping
+MLX_TO_TORCH_DTYPE: dict[mx.Dtype, torch.dtype] = {
+    mx.float32: torch.float32,
+    mx.float16: torch.float16,
+    mx.bfloat16: torch.bfloat16,
+    mx.int32: torch.int32,
+    mx.int64: torch.int64,
+    mx.int16: torch.int16,
+    mx.int8: torch.int8,
+    mx.uint8: torch.uint8,
+    mx.bool_: torch.bool,
+}
+
+# PyTorch to MLX dtype mapping
+TORCH_TO_MLX_DTYPE: dict[torch.dtype, mx.Dtype] = {
+    v: k for k, v in MLX_TO_TORCH_DTYPE.items()
+}
+
+
+def get_torch_device() -> torch.device:
+    """Get the PyTorch device for Metal/MPS.
+
+    Returns:
+        torch.device for MPS if available, else CPU
+    """
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def _get_tensor_size_bytes(array: mx.array) -> int:
+    """Calculate the size of an MLX array in bytes.
+
+    Args:
+        array: MLX array
+
+    Returns:
+        Size in bytes
+    """
+    return array.size * array.dtype.size
+
+
+def _is_safe_for_mps(array: mx.array) -> bool:
+    """Check if an array is safe to transfer to MPS without hitting size limits.
+
+    MPS has a 4GB limit for MPSTemporaryNDArray, but Metal may allocate
+    multiple temporary buffers internally. We use a conservative threshold.
+
+    Args:
+        array: MLX array to check
+
+    Returns:
+        True if safe to transfer to MPS, False if should stay on CPU
+    """
+    return _get_tensor_size_bytes(array) < _MPS_SAFE_SIZE_BYTES
+
+
+def torch_to_mlx(tensor: torch.Tensor) -> mx.array:
+    """Convert PyTorch tensor to MLX array.
+
+    Uses numpy as an intermediate to enable zero-copy on unified memory.
+
+    Args:
+        tensor: PyTorch tensor (can be on any device)
+
+    Returns:
+        MLX array with the same data
+    """
+    # Move to CPU if on MPS for numpy conversion
+    if tensor.device.type != "cpu":
+        tensor = tensor.cpu()
+
+    tensor = tensor.detach()
+
+    # Note: numpy does not support bfloat16.
+    if tensor.dtype == torch.bfloat16:
+        return mx.array(tensor)
+
+    return mx.array(tensor.numpy())
+
+
+def mlx_to_torch(
+    array: mx.array,
+    device: torch.device | Literal["mps", "cpu"] | None = None,
+    already_contiguous: bool = False,
+) -> torch.Tensor:
+    """Convert MLX array to PyTorch tensor.
+
+    Uses numpy as an intermediate to enable zero-copy on unified memory.
+
+    Args:
+        array: MLX array
+        device: Target PyTorch device (default: MPS if available)
+        already_contiguous: Skip contiguity check if array is known contiguous
+
+    Returns:
+        PyTorch tensor with the same data
+    """
+    if device is None:
+        device = get_torch_device()
+    elif isinstance(device, str):
+        device = torch.device(device)
+
+    # Use memoryview for zero-copy conversion (bypasses numpy for bfloat16)
+    # reference: https://github.com/ml-explore/mlx/issues/403
+    torch_dtype = MLX_TO_TORCH_DTYPE.get(array.dtype)
+    if torch_dtype is not None:
+        if already_contiguous:
+            # Fast path: skip contiguity check, single eval
+            mx.eval(array)
+            buffer = memoryview(array)
+        else:
+            # MLX views / non-contiguous arrays expose a non-contiguous buffer (or
+            # sometimes no usable buffer), which `torch.frombuffer` can't consume.
+            # Make contiguous first, then eval once
+            array = mx.contiguous(array)
+            mx.eval(array)
+            buffer = memoryview(array)
+
+        tensor = torch.frombuffer(buffer, dtype=torch_dtype).reshape(array.shape)
+    else:
+        # Fallback to numpy path for unsupported dtypes
+        raise ValueError(f"Unsupported MLX dtype: {array.dtype}")
+
+    # Move to target device, but check for MPS size limits first
+    if device.type == "mps":
+        # Ensure all MLX Metal commands complete before MPS uses the GPU
+        sync_mlx()
+        if _is_safe_for_mps(array):
+            tensor = tensor.to(device)
+        else:
+            # Large tensor - keep on CPU to avoid MPS 4GB limit crash
+            # See: https://github.com/anthropics/vllm-metal/issues/43
+            logger.debug(
+                "Tensor too large for MPS (%d bytes > %d limit), keeping on CPU",
+                _get_tensor_size_bytes(array),
+                _MPS_SAFE_SIZE_BYTES,
+            )
+    elif device.type != "cpu":
+        tensor = tensor.to(device)
+
+    return tensor
+
+
+def sync_mlx() -> None:
+    """Synchronize MLX operations.
+
+    Call this before converting MLX arrays to ensure all operations complete.
+    """
+    # Prefer an explicit MLX barrier when available; otherwise force evaluation.
+    # `mx.eval([])` is a no-op, so we evaluate a tiny scalar as a safe fallback.
+    try:
+        mx.synchronize()
+    except (AttributeError, TypeError):
+        mx.eval(mx.array(0, dtype=mx.int32))
+
+
+def sync_torch() -> None:
+    """Synchronize PyTorch MPS operations.
+
+    Call this before converting PyTorch tensors to ensure all operations complete.
+    """
+    if torch.backends.mps.is_available():
+        torch.mps.synchronize()

mlx_tinygrad_interop/my_interop.py

@@ -0,0 +1,34 @@
+from tinygrad.helpers import GlobalCounters
+from tinygrad.tensor import Tensor
+from tinygrad.device import Device
+from tinygrad.dtype import _from_torch_dtype, _from_np_dtype
+import torch
+import time
+import statistics
+
+def main() -> None:
+    for i in range(8):
+        N = 256 * (4 ** i)
+        x = torch.zeros(N, device=torch.device("mps"))
+        
+        vals = []
+        for j in range(1000):
+            x = x.uniform_()
+            torch.mps.synchronize()
+
+            old = time.perf_counter_ns()
+            Tensor.from_blob(x.data_ptr(), x.shape, dtype=_from_torch_dtype(x.dtype), device="METAL")
+            Tensor.from_blob(x.data_ptr(), x.shape, dtype=_from_torch_dtype(x.dtype), device="METAL")
+            Tensor.from_blob(x.data_ptr(), x.shape, dtype=_from_torch_dtype(x.dtype), device="METAL")
+            Tensor.from_blob(x.data_ptr(), x.shape, dtype=_from_torch_dtype(x.dtype), device="METAL")
+            Tensor.from_blob(x.data_ptr(), x.shape, dtype=_from_torch_dtype(x.dtype), device="METAL")
+            Tensor.from_blob(x.data_ptr(), x.shape, dtype=_from_torch_dtype(x.dtype), device="METAL")
+            Tensor.from_blob(x.data_ptr(), x.shape, dtype=_from_torch_dtype(x.dtype), device="METAL")
+            Tensor.from_blob(x.data_ptr(), x.shape, dtype=_from_torch_dtype(x.dtype), device="METAL")
+            new = time.perf_counter_ns()
+            vals.append(float(new - old) / 8)
+        print( f"result: {N*4:8d} pytorch to tinygrad in {statistics.mean(vals):.2f}ns with stddev {statistics.stdev(vals):.2f}ns" )
+
+
+if __name__=="__main__":
+    main()

mlx_tinygrad_interop/stress_interop.py

@@ -0,0 +1,473 @@
+from __future__ import annotations
+
+import argparse
+import gc
+import resource
+import sys
+import tracemalloc
+from typing import Any, cast
+
+import mlx.core as mx
+import numpy as np
+from tinygrad import Device, Tensor, dtypes
+from tinygrad.device import Buffer
+
+try:
+  from mlx_tinygrad_interop.lib.lease_pool import MlxToTinygradCopyLeasePools, MlxToTinygradLeasePools
+except ModuleNotFoundError:
+  from lib.lease_pool import MlxToTinygradCopyLeasePools, MlxToTinygradLeasePools
+
+DTYPES: dict[str, tuple[Any, Any, np.dtype[Any]]] = {
+  "float16": (mx.float16, dtypes.float16, np.dtype(np.float16)),
+  "float32": (mx.float32, dtypes.float32, np.dtype(np.float32)),
+  "int32": (mx.int32, dtypes.int32, np.dtype(np.int32)),
+  "uint8": (mx.uint8, dtypes.uint8, np.dtype(np.uint8)),
+}
+
+
+def parse_args() -> argparse.Namespace:
+  parser = argparse.ArgumentParser(description="Randomized correctness and soak checks for MLX <-> tinygrad interop.")
+  parser.add_argument("--seed", type=int, default=0)
+  parser.add_argument("--cases", type=int, default=64, help="Number of randomized correctness cases.")
+  parser.add_argument("--soak-iterations", type=int, default=512, help="Number of repeated pool/copy iterations after correctness checks.")
+  parser.add_argument("--max-elements", type=int, default=4096, help="Upper bound on random tensor element count.")
+  parser.add_argument("--max-pools", type=int, default=16, help="Maximum keyed pools to retain for alias and copy handoff managers.")
+  parser.add_argument("--dtypes", default="float16,float32,int32,uint8", help="Comma-separated dtype names to exercise.")
+  return parser.parse_args()
+
+
+def tinygrad_zero_copy_memoryview(t: Tensor) -> memoryview:
+  buf = cast(Buffer, t.uop.buffer).ensure_allocated()
+  assert t.dtype.base.fmt is not None, f"no buffer format for dtype {t.dtype.base}"
+  return buf.as_memoryview(force_zero_copy=True).cast(t.dtype.base.fmt, t.shape)
+
+
+def tinygrad_from_mlx_fast(x: Any, tg_dtype: Any) -> Tensor:
+  storage = mx.metal._unsafe_export_storage(x)
+  return Tensor._unsafe_from_metal_buffer_fast(
+    int(storage["mtl_buffer_ptr"]),
+    tuple(storage["shape"]),
+    dtype=tg_dtype,
+    byte_offset=int(storage["offset_bytes"]),
+    buffer_nbytes=int(storage["buffer_nbytes"]),
+    owner=x,
+  )
+
+
+def tinygrad_from_mlx_single_entry(x: Any, tg_dtype: Any) -> Tensor:
+  return mx.metal._unsafe_to_tinygrad_fast(x, tg_dtype, owner=x)
+
+
+def mlx_from_tinygrad_alias_only(t: Tensor, mx_dtype: Any) -> Any:
+  storage = t._unsafe_metal_storage()
+  return mx.metal._unsafe_array_from_ptr_alias_only(
+    int(storage["raw_ptr"]),
+    tuple(storage["shape"]),
+    mx_dtype,
+    owner=t,
+  )
+
+
+def mlx_from_tinygrad_maybe_copy(t: Tensor, mx_dtype: Any) -> Any:
+  storage = t._unsafe_metal_storage()
+  return mx.metal._unsafe_array_from_ptr(
+    int(storage["raw_ptr"]),
+    tuple(storage["shape"]),
+    mx_dtype,
+    owner=t,
+  )
+
+
+def mlx_from_tinygrad_copy(t: Tensor) -> Any:
+  return mx.array(tinygrad_zero_copy_memoryview(t))
+
+
+def assert_array_close(name: str, actual: np.ndarray, expected: np.ndarray) -> None:
+  if np.issubdtype(expected.dtype, np.floating):
+    if expected.dtype == np.float16:
+      rtol, atol = 5e-3, 5e-3
+    else:
+      # Mixed matmul/reduction chains across NumPy/MLX/tinygrad can drift by a
+      # few ulps from accumulation-order differences even when the conversion is
+      # correct. Keep float32 strict, but not unrealistically bit-exact.
+      rtol, atol = 5e-5, 1e-5
+    np.testing.assert_allclose(actual, expected, rtol=rtol, atol=atol, err_msg=name)
+  else:
+    np.testing.assert_array_equal(actual, expected, err_msg=name)
+
+
+def rss_max_bytes() -> int:
+  rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+  return int(rss if sys.platform == "darwin" else rss * 1024)
+
+
+def mlx_memory_snapshot() -> dict[str, int]:
+  stats: dict[str, int] = {}
+  for name in ("get_active_memory", "get_cache_memory", "get_peak_memory"):
+    fn = getattr(mx, name, None)
+    if callable(fn): stats[name] = int(fn())
+  return stats
+
+
+def random_shape(rng: np.random.Generator, max_elements: int) -> tuple[int, ...]:
+  ndim = int(rng.integers(1, 5))
+  remaining = max(1, int(rng.integers(1, max_elements + 1)))
+  shape: list[int] = []
+  for dim_index in range(ndim):
+    dims_left = ndim - dim_index
+    if dims_left == 1:
+      shape.append(remaining)
+      break
+    dim = int(rng.integers(1, max(2, int(round(remaining ** (1 / dims_left))) + 2)))
+    shape.append(dim)
+    remaining = max(1, remaining // dim)
+  return tuple(shape)
+
+
+def random_values(rng: np.random.Generator, np_dtype: np.dtype[Any], shape: tuple[int, ...]) -> np.ndarray:
+  if np.issubdtype(np_dtype, np.floating):
+    data = rng.standard_normal(np.prod(shape, dtype=np.int64)).astype(np.float32)
+    return data.astype(np_dtype).reshape(shape)
+  if np.issubdtype(np_dtype, np.unsignedinteger):
+    info = np.iinfo(np_dtype)
+    return rng.integers(0, min(info.max, 255) + 1, size=shape, dtype=np_dtype)
+  info = np.iinfo(np_dtype)
+  return rng.integers(max(info.min, -128), min(info.max, 127) + 1, size=shape, dtype=np_dtype)
+
+
+def make_mlx_source(values: np.ndarray, mx_dtype: Any, rng: np.random.Generator) -> Any:
+  flat = values.reshape(-1)
+  if flat.size > 0 and rng.random() < 0.5:
+    offset_elems = int(rng.integers(1, 33))
+    suffix = int(rng.integers(1, 33))
+    backing = np.zeros(offset_elems + flat.size + suffix, dtype=flat.dtype)
+    backing[offset_elems:offset_elems + flat.size] = flat
+    base = mx.array(backing, dtype=mx_dtype)
+    view = base[offset_elems:offset_elems + flat.size].reshape(values.shape)
+    mx.eval(view)
+    return view
+  return mx.array(values, dtype=mx_dtype)
+
+
+def random_scalar(rng: np.random.Generator, np_dtype: np.dtype[Any]) -> Any:
+  if np.issubdtype(np_dtype, np.floating):
+    return np_dtype.type(rng.uniform(-2.0, 2.0)).item()
+  if np.issubdtype(np_dtype, np.unsignedinteger):
+    return int(rng.integers(0, 4))
+  return int(rng.integers(-3, 4))
+
+
+def random_slice(axis_size: int, rng: np.random.Generator) -> tuple[int, int]:
+  start = int(rng.integers(0, axis_size))
+  end = int(rng.integers(start + 1, axis_size + 1))
+  return start, end
+
+
+def random_ops(rng: np.random.Generator, shape: tuple[int, ...], np_dtype: np.dtype[Any]) -> list[tuple[str, Any]]:
+  ops: list[tuple[str, Any]] = []
+  cur_shape = shape
+  if len(cur_shape) > 1 and rng.random() < 0.7:
+    perm = tuple(int(x) for x in rng.permutation(len(cur_shape)))
+    ops.append(("transpose", perm))
+    cur_shape = tuple(cur_shape[i] for i in perm)
+  if len(cur_shape) > 1 and rng.random() < 0.6:
+    reshaped = tuple(reversed(cur_shape))
+    ops.append(("reshape", reshaped))
+    cur_shape = reshaped
+  if any(dim > 1 for dim in cur_shape) and rng.random() < 0.5:
+    axis = int(rng.choice([i for i, dim in enumerate(cur_shape) if dim > 1]))
+    start, end = random_slice(cur_shape[axis], rng)
+    ops.append(("slice", (axis, start, end)))
+    cur_shape = cur_shape[:axis] + (end - start,) + cur_shape[axis + 1:]
+  ops.append(("add", random_scalar(rng, np_dtype)))
+  ops.append(("mul", random_scalar(rng, np_dtype)))
+  if cur_shape and rng.random() < 0.7:
+    bshape = tuple(dim if rng.random() < 0.5 else 1 for dim in cur_shape)
+    ops.append(("broadcast_add", random_values(rng, np_dtype, bshape)))
+  if rng.random() < 0.5:
+    ops.append(("relu", None))
+  if np.issubdtype(np_dtype, np.floating) and cur_shape and rng.random() < 0.35:
+    out_cols = int(rng.integers(1, min(8, cur_shape[-1]) + 1))
+    weight = random_values(rng, np_dtype, (cur_shape[-1], out_cols))
+    ops.append(("matmul_lastdim", weight))
+    cur_shape = (int(np.prod(cur_shape[:-1], dtype=np.int64)), out_cols)
+  if cur_shape and rng.random() < 0.5:
+    axis = int(rng.integers(0, len(cur_shape)))
+    keepdim = bool(rng.integers(0, 2))
+    ops.append(("sum", (axis, keepdim)))
+    cur_shape = cur_shape[:axis] + ((1,) if keepdim else ()) + cur_shape[axis + 1:]
+  if cur_shape and rng.random() < 0.35:
+    axis = int(rng.integers(0, len(cur_shape)))
+    ops.append(("concat_self", axis))
+  return ops
+
+
+def _slice_spec(shape: tuple[int, ...], axis: int, start: int, end: int) -> tuple[slice, ...]:
+  return tuple(slice(start, end) if i == axis else slice(None) for i in range(len(shape)))
+
+
+def apply_numpy_ops(x: np.ndarray, ops: list[tuple[str, Any]]) -> np.ndarray:
+  out = x.copy()
+  for op, arg in ops:
+    if op == "transpose":
+      out = np.transpose(out, arg)
+    elif op == "reshape":
+      out = out.reshape(arg)
+    elif op == "slice":
+      axis, start, end = arg
+      out = out[_slice_spec(out.shape, axis, start, end)]
+    elif op == "add":
+      out = out + arg
+    elif op == "mul":
+      out = out * arg
+    elif op == "broadcast_add":
+      out = out + arg
+    elif op == "relu":
+      out = np.maximum(out, 0)
+    elif op == "matmul_lastdim":
+      # NumPy's `@` path on the current macOS validation host produced an
+      # incorrect all-zero result for a valid contiguous float32 case that MLX,
+      # tinygrad, and `np.einsum` all agreed on. Use einsum here so the stress
+      # harness keeps a trustworthy numerical baseline.
+      out = np.einsum("ik,kj->ij", out.reshape(-1, out.shape[-1]), arg, optimize=True)
+    elif op == "sum":
+      axis, keepdim = arg
+      out = out.sum(axis=axis, keepdims=keepdim)
+    elif op == "concat_self":
+      out = np.concatenate([out, out], axis=arg)
+    else:
+      raise RuntimeError(f"unknown op {op}")
+  return out
+
+
+def apply_tinygrad_ops(x: Tensor, ops: list[tuple[str, Any]]) -> Tensor:
+  out = x
+  for op, arg in ops:
+    if op == "transpose":
+      out = out.permute(arg)
+    elif op == "reshape":
+      out = out.reshape(arg)
+    elif op == "slice":
+      axis, start, end = arg
+      out = out[_slice_spec(out.shape, axis, start, end)]
+    elif op == "add":
+      out = out + arg
+    elif op == "mul":
+      out = out * arg
+    elif op == "broadcast_add":
+      out = out + Tensor(arg, device=out.device, dtype=out.dtype)
+    elif op == "relu":
+      out = out.relu()
+    elif op == "matmul_lastdim":
+      out = out.reshape(-1, out.shape[-1]) @ Tensor(arg, device=out.device, dtype=out.dtype)
+    elif op == "sum":
+      axis, keepdim = arg
+      out = out.sum(axis=axis, keepdim=keepdim)
+    elif op == "concat_self":
+      out = out.cat(out, dim=arg)
+    else:
+      raise RuntimeError(f"unknown op {op}")
+  return out.realize()
+
+
+def apply_mlx_ops(x: Any, ops: list[tuple[str, Any]]) -> Any:
+  out = x
+  for op, arg in ops:
+    if op == "transpose":
+      out = mx.transpose(out, arg)
+    elif op == "reshape":
+      out = mx.reshape(out, arg)
+    elif op == "slice":
+      axis, start, end = arg
+      out = out[_slice_spec(out.shape, axis, start, end)]
+    elif op == "add":
+      out = out + arg
+    elif op == "mul":
+      out = out * arg
+    elif op == "broadcast_add":
+      out = out + mx.array(arg, dtype=out.dtype)
+    elif op == "relu":
+      out = mx.maximum(out, 0)
+    elif op == "matmul_lastdim":
+      out = mx.reshape(out, (-1, out.shape[-1])) @ mx.array(arg, dtype=out.dtype)
+    elif op == "sum":
+      axis, keepdim = arg
+      out = mx.sum(out, axis=axis, keepdims=keepdim)
+    elif op == "concat_self":
+      out = mx.concatenate([out, out], axis=arg)
+    else:
+      raise RuntimeError(f"unknown op {op}")
+  mx.eval(out)
+  return out
+
+
+def run_case(case_index: int, rng: np.random.Generator, mx_dtype: Any, tg_dtype: Any, np_dtype: np.dtype[Any],
+             max_elements: int, alias_pools: MlxToTinygradLeasePools, copy_pools: MlxToTinygradCopyLeasePools) -> None:
+  shape = random_shape(rng, max_elements)
+  values = random_values(rng, np_dtype, shape)
+  ops = random_ops(rng, shape, np_dtype)
+
+  mx_source = make_mlx_source(values, mx_dtype, rng)
+  tg_source = Tensor(values, device="METAL", dtype=tg_dtype).realize()
+  mx_baseline_source = mx.array(values, dtype=mx_dtype)
+  Device["METAL"].synchronize()
+
+  expected_raw = values
+  tg_expected_after_ops = apply_tinygrad_ops(tg_source, ops).numpy()
+  mlx_expected_after_ops = np.array(apply_mlx_ops(mx_baseline_source, ops))
+  tg_fast = tinygrad_from_mlx_fast(mx_source, tg_dtype)
+  tg_single = tinygrad_from_mlx_single_entry(mx_source, tg_dtype)
+  assert_array_close(f"case {case_index} mlx->tinygrad fast raw", tg_fast.numpy(), expected_raw)
+  assert_array_close(f"case {case_index} mlx->tinygrad single raw", tg_single.numpy(), expected_raw)
+  assert_array_close(
+    f"case {case_index} mlx->tinygrad fast ops",
+    apply_tinygrad_ops(tg_fast, ops).numpy(),
+    tg_expected_after_ops,
+  )
+  assert_array_close(
+    f"case {case_index} mlx->tinygrad single ops",
+    apply_tinygrad_ops(tg_single, ops).numpy(),
+    tg_expected_after_ops,
+  )
+
+  alias_result = alias_pools.run_with_mlx_tensor(
+    mx_source,
+    tg_dtype=tg_dtype,
+    fn=lambda tg: apply_tinygrad_ops(tg, ops),
+  )
+  copy_result = copy_pools.run_with_mlx_tensor(
+    mx_source,
+    tg_dtype=tg_dtype,
+    fn=lambda tg: apply_tinygrad_ops(tg, ops),
+  )
+  assert_array_close(f"case {case_index} alias scoped ops", alias_result.numpy(), tg_expected_after_ops)
+  assert_array_close(f"case {case_index} copy scoped ops", copy_result.numpy(), tg_expected_after_ops)
+
+  alias_roundtrip = np.array(mlx_from_tinygrad_copy(alias_result.cast(tg_dtype).realize()))
+  copy_roundtrip = np.array(mlx_from_tinygrad_copy(copy_result.cast(tg_dtype).realize()))
+  assert_array_close(f"case {case_index} alias roundtrip raw", alias_roundtrip, np.array(alias_result.cast(tg_dtype).numpy(), copy=True))
+  assert_array_close(f"case {case_index} copy roundtrip raw", copy_roundtrip, np.array(copy_result.cast(tg_dtype).numpy(), copy=True))
+
+  mx_alias = mlx_from_tinygrad_alias_only(tg_source, mx_dtype)
+  mx_maybe_copy = mlx_from_tinygrad_maybe_copy(tg_source, mx_dtype)
+  mx_copy = mlx_from_tinygrad_copy(tg_source)
+  assert_array_close(f"case {case_index} tinygrad->mlx alias raw", np.array(mx_alias), expected_raw)
+  assert_array_close(f"case {case_index} tinygrad->mlx maybe_copy raw", np.array(mx_maybe_copy), expected_raw)
+  assert_array_close(f"case {case_index} tinygrad->mlx copy raw", np.array(mx_copy), expected_raw)
+  assert_array_close(
+    f"case {case_index} tinygrad->mlx alias ops",
+    np.array(apply_mlx_ops(mx_alias, ops)),
+    mlx_expected_after_ops,
+  )
+  assert_array_close(
+    f"case {case_index} tinygrad->mlx maybe_copy ops",
+    np.array(apply_mlx_ops(mx_maybe_copy, ops)),
+    mlx_expected_after_ops,
+  )
+  assert_array_close(
+    f"case {case_index} tinygrad->mlx copy ops",
+    np.array(apply_mlx_ops(mx_copy, ops)),
+    mlx_expected_after_ops,
+  )
+
+
+def run_soak(rng: np.random.Generator, dtype_names: list[str], iterations: int, max_elements: int, max_pools: int) -> tuple[float, list[tuple[int, dict[str, int], int, int, int]]]:
+  alias_pools = MlxToTinygradLeasePools(capacity_per_key=8, max_pools=max_pools, synchronize_on_release=True)
+  copy_pools = MlxToTinygradCopyLeasePools(capacity_per_key=8, max_pools=max_pools, synchronize_on_release=True)
+  checksum = 0.0
+  native_checkpoints: list[tuple[int, dict[str, int], int, int, int]] = []
+  for iteration in range(iterations):
+    dtype_name = dtype_names[iteration % len(dtype_names)]
+    mx_dtype, tg_dtype, np_dtype = DTYPES[dtype_name]
+    shape = random_shape(rng, max_elements)
+    values = random_values(rng, np_dtype, shape)
+    mx_source = make_mlx_source(values, mx_dtype, rng)
+    tg_source = Tensor(values, device="METAL", dtype=tg_dtype).realize()
+    Device["METAL"].synchronize()
+
+    checksum += float(alias_pools.run_with_mlx_tensor(
+      mx_source,
+      tg_dtype=tg_dtype,
+      fn=lambda tg: apply_tinygrad_ops(tg, [("add", 1.0 if np.issubdtype(np_dtype, np.floating) else 1)]).sum(),
+    ).item())
+    checksum += float(copy_pools.run_with_mlx_tensor(
+      mx_source,
+      tg_dtype=tg_dtype,
+      fn=lambda tg: apply_tinygrad_ops(tg, [("mul", 1.0 if np.issubdtype(np_dtype, np.floating) else 1)]).sum(),
+    ).item())
+    checksum += float(np.array(mlx_from_tinygrad_copy(tg_source)).astype(np.float64).sum())
+    if (iteration + 1) % 64 == 0:
+      gc.collect()
+      native_checkpoints.append((iteration + 1, mlx_memory_snapshot(), rss_max_bytes(), alias_pools.pool_count, copy_pools.pool_count))
+
+  if alias_pools.pool_count > max_pools:
+    raise AssertionError(f"alias pool registry exceeded cap: {alias_pools.pool_count} > {max_pools}")
+  if copy_pools.pool_count > max_pools:
+    raise AssertionError(f"copy pool registry exceeded cap: {copy_pools.pool_count} > {max_pools}")
+  print(f"# soak_checksum={checksum:.6f} alias_pool_count={alias_pools.pool_count} copy_pool_count={copy_pools.pool_count}")
+  return checksum, native_checkpoints
+
+
+def main() -> None:
+  args = parse_args()
+  if not mx.metal.is_available():
+    raise RuntimeError("Metal is not available")
+
+  dtype_names = [name.strip() for name in args.dtypes.split(",") if name.strip()]
+  missing = [name for name in dtype_names if name not in DTYPES]
+  if missing:
+    raise RuntimeError(f"unknown dtype(s): {', '.join(missing)}")
+
+  rng = np.random.default_rng(args.seed)
+  tracemalloc.start()
+  before_cur, before_peak = tracemalloc.get_traced_memory()
+  reset_peak = getattr(mx, "reset_peak_memory", None)
+  if callable(reset_peak): reset_peak()
+  native_before = mlx_memory_snapshot()
+  rss_before = rss_max_bytes()
+  alias_pools = MlxToTinygradLeasePools(capacity_per_key=8, max_pools=args.max_pools, synchronize_on_release=True)
+  copy_pools = MlxToTinygradCopyLeasePools(capacity_per_key=8, max_pools=args.max_pools, synchronize_on_release=True)
+
+  for case_index in range(args.cases):
+    dtype_name = dtype_names[case_index % len(dtype_names)]
+    run_case(case_index, rng, *DTYPES[dtype_name], max_elements=args.max_elements, alias_pools=alias_pools, copy_pools=copy_pools)
+
+  _, native_checkpoints = run_soak(rng, dtype_names, iterations=args.soak_iterations, max_elements=args.max_elements, max_pools=args.max_pools)
+  gc.collect()
+  after_cur, after_peak = tracemalloc.get_traced_memory()
+  native_after = mlx_memory_snapshot()
+  rss_after = rss_max_bytes()
+  active_peak = max([native_before.get("get_active_memory", 0), native_after.get("get_active_memory", 0),
+                     *[stats.get("get_active_memory", 0) for _, stats, _, _, _ in native_checkpoints]], default=0)
+  cache_peak = max([native_before.get("get_cache_memory", 0), native_after.get("get_cache_memory", 0),
+                    *[stats.get("get_cache_memory", 0) for _, stats, _, _, _ in native_checkpoints]], default=0)
+  rss_peak = max([rss_before, rss_after, *[rss for _, _, rss, _, _ in native_checkpoints]], default=0)
+  alias_pool_peak = max([alias_pools.pool_count, *[alias_pool_count for _, _, _, alias_pool_count, _ in native_checkpoints]], default=0)
+  copy_pool_peak = max([copy_pools.pool_count, *[copy_pool_count for _, _, _, _, copy_pool_count in native_checkpoints]], default=0)
+  print(
+    "# native_memory"
+    f" mlx_active_start={native_before.get('get_active_memory', -1)}"
+    f" mlx_active_end={native_after.get('get_active_memory', -1)}"
+    f" mlx_active_peak={active_peak}"
+    f" mlx_cache_start={native_before.get('get_cache_memory', -1)}"
+    f" mlx_cache_end={native_after.get('get_cache_memory', -1)}"
+    f" mlx_cache_peak={cache_peak}"
+    f" mlx_reported_peak={native_after.get('get_peak_memory', -1)}"
+    f" rss_max_start={rss_before}"
+    f" rss_max_end={rss_after}"
+    f" rss_max_peak={rss_peak}"
+    f" alias_pool_peak={alias_pool_peak}"
+    f" copy_pool_peak={copy_pool_peak}"
+  )
+  print(
+    "# stress_ok"
+    f" cases={args.cases}"
+    f" soak_iterations={args.soak_iterations}"
+    f" max_pools={args.max_pools}"
+    f" tracemalloc_current={after_cur - before_cur}"
+    f" tracemalloc_peak={after_peak - before_peak}"
+  )
+
+
+if __name__ == "__main__":
+  main()

mlx_tinygrad_interop/tensor_bridge.py

@@ -0,0 +1,3 @@
+"""Compatibility shim for the moved tensor bridge implementation."""
+
+from mlx_tinygrad_interop.lib.tensor_bridge import *  # noqa: F401,F403

mlx_tinygrad_interop/tensor_bridge_vllm.py

@@ -0,0 +1,3 @@
+"""Compatibility shim for the moved vLLM-style torch bridge module."""
+
+from mlx_tinygrad_interop.lib.tensor_bridge_vllm import *  # noqa: F401,F403

mlx_tinygrad_interop/test_handoff.py

@@ -0,0 +1,221 @@
+import unittest
+
+import mlx.core as mx
+import numpy as np
+from tinygrad import Tensor, dtypes
+
+from mlx_tinygrad_interop.lib.lease_pool import (
+  MlxToTinygradCopyKey,
+  MlxToTinygradCopyLeasePool,
+  MlxToTinygradCopyLeasePools,
+  MlxToTinygradLeaseKey,
+  MlxToTinygradLeasePool,
+  MlxToTinygradLeasePools,
+)
+
+
+@unittest.skipUnless(mx.metal.is_available(), "Metal is not available")
+class TestMlxTinygradLeasePool(unittest.TestCase):
+  def test_checked_rebind_requires_shape_and_dtype(self):
+    array = mx.array(np.arange(16, dtype=np.float32), dtype=mx.float32)
+    storage = mx.metal._unsafe_export_storage(array)
+    borrower = Tensor._unsafe_metal_borrower(
+      int(storage["mtl_buffer_ptr"]),
+      tuple(storage["shape"]),
+      dtype=dtypes.float32,
+      byte_offset=int(storage["offset_bytes"]),
+      buffer_nbytes=int(storage["buffer_nbytes"]),
+      owner=array,
+    )
+
+    with self.assertRaises(TypeError):
+      borrower.rebind(
+        int(storage["mtl_buffer_ptr"]),
+        owner=array,
+        byte_offset=int(storage["offset_bytes"]),
+        buffer_nbytes=int(storage["buffer_nbytes"]),
+      )
+
+  def test_lease_pool_evaluates_lazy_array_on_acquire(self):
+    lazy = mx.arange(32, dtype=mx.float32) + 3
+    pool = MlxToTinygradLeasePool.from_mlx(lazy, tg_dtype=dtypes.float32, capacity=1)
+
+    with pool.acquire_from_mlx(lazy) as lease:
+      np.testing.assert_array_equal(lease.tensor.numpy(), np.arange(32, dtype=np.float32) + np.float32(3))
+
+  def test_raw_lease_tensor_can_observe_reused_slot_contents(self):
+    first = mx.array(np.arange(16, dtype=np.float32), dtype=mx.float32)
+    second = mx.array(np.arange(16, dtype=np.float32) + np.float32(10), dtype=mx.float32)
+    pool = MlxToTinygradLeasePool.from_mlx(first, tg_dtype=dtypes.float32, capacity=1)
+
+    lease_a = pool.acquire_from_mlx(first)
+    raw_tensor = lease_a.tensor
+    lease_a.release()
+
+    with pool.acquire_from_mlx(second):
+      np.testing.assert_array_equal(raw_tensor.numpy(), np.arange(16, dtype=np.float32) + np.float32(10))
+
+  def test_lease_pool_capacity_requires_release(self):
+    first = mx.array(np.arange(16, dtype=np.float32), dtype=mx.float32)
+    second = mx.array(np.arange(16, dtype=np.float32) + np.float32(1), dtype=mx.float32)
+    pool = MlxToTinygradLeasePool.from_mlx(first, tg_dtype=dtypes.float32, capacity=1)
+
+    lease = pool.acquire_from_mlx(first)
+    with self.assertRaisesRegex(RuntimeError, "still in use"):
+      pool.acquire_from_mlx(second)
+    lease.release(synchronize=False)
+
+    with pool.acquire_from_mlx(second) as second_lease:
+      np.testing.assert_array_equal(second_lease.tensor.numpy(), np.arange(16, dtype=np.float32) + np.float32(1))
+
+  def test_lease_release_invalidates_future_access(self):
+    array = mx.array(np.arange(8, dtype=np.float32), dtype=mx.float32)
+    pool = MlxToTinygradLeasePool.from_mlx(array, tg_dtype=dtypes.float32, capacity=1)
+
+    lease = pool.acquire_from_mlx(array)
+    self.assertEqual(lease.generation, 1)
+    lease.release(synchronize=False)
+
+    with self.assertRaisesRegex(RuntimeError, "already released"):
+      _ = lease.tensor
+    with self.assertRaisesRegex(RuntimeError, "already released"):
+      lease.release(synchronize=False)
+
+  def test_safe_release_clears_owner_reference(self):
+    array = mx.array(np.arange(8, dtype=np.float32), dtype=mx.float32)
+    pool = MlxToTinygradLeasePool.from_mlx(array, tg_dtype=dtypes.float32, capacity=1)
+
+    lease = pool.acquire_from_mlx(array)
+    self.assertTrue(hasattr(pool._slots[0].borrower._base_buf, "_external_owner"))
+    lease.release()
+    self.assertFalse(hasattr(pool._slots[0].borrower._base_buf, "_external_owner"))
+
+  def test_unsafe_release_without_sync_retains_owner_reference(self):
+    array = mx.array(np.arange(8, dtype=np.float32), dtype=mx.float32)
+    pool = MlxToTinygradLeasePool.from_mlx(array, tg_dtype=dtypes.float32, capacity=1, synchronize_on_release=False)
+
+    lease = pool.acquire_from_mlx(array)
+    lease.release(synchronize=False)
+    self.assertTrue(hasattr(pool._slots[0].borrower._base_buf, "_external_owner"))
+
+  def test_scoped_handoff_realizes_returned_tensor_before_release(self):
+    first = mx.array(np.arange(16, dtype=np.float32), dtype=mx.float32)
+    second = mx.array(np.arange(16, dtype=np.float32) + np.float32(5), dtype=mx.float32)
+    pools = MlxToTinygradLeasePools(capacity_per_key=1, synchronize_on_release=True)
+
+    result_first = pools.run_with_mlx_tensor(first, tg_dtype=dtypes.float32, fn=lambda t: (t + 1).sum())
+    result_second = pools.run_with_mlx_tensor(second, tg_dtype=dtypes.float32, fn=lambda t: (t + 1).sum())
+
+    np.testing.assert_allclose(result_first.numpy(), np.array((np.arange(16, dtype=np.float32) + 1).sum(), dtype=np.float32))
+    np.testing.assert_allclose(result_second.numpy(), np.array((np.arange(16, dtype=np.float32) + 6).sum(), dtype=np.float32))
+
+  def test_scoped_handoff_rejects_returning_alias_view(self):
+    array = mx.array(np.arange(16, dtype=np.float32), dtype=mx.float32)
+    pools = MlxToTinygradLeasePools(capacity_per_key=1, synchronize_on_release=True)
+
+    with self.assertRaisesRegex(RuntimeError, "alias the borrowed slot"):
+      pools.run_with_mlx_tensor(array, tg_dtype=dtypes.float32, fn=lambda t: t.reshape(2, 8))
+
+  def test_scoped_handoff_rejects_stashed_lazy_tensor(self):
+    array = mx.array(np.arange(16, dtype=np.float32), dtype=mx.float32)
+    pools = MlxToTinygradLeasePools(capacity_per_key=1, synchronize_on_release=True)
+    stash: dict[str, Tensor] = {}
+
+    with self.assertRaisesRegex(RuntimeError, "leaked tensor\\(s\\) derived"):
+      pools.run_with_mlx_tensor(
+        array,
+        tg_dtype=dtypes.float32,
+        fn=lambda t: (stash.setdefault("u", t + 1), (t + 2).sum())[1],
+      )
+
+  def test_scoped_handoff_rejects_stashed_alias_view(self):
+    array = mx.array(np.arange(16, dtype=np.float32), dtype=mx.float32)
+    pools = MlxToTinygradLeasePools(capacity_per_key=1, synchronize_on_release=True)
+    stash: dict[str, Tensor] = {}
+
+    with self.assertRaisesRegex(RuntimeError, "leaked tensor\\(s\\) derived"):
+      pools.run_with_mlx_tensor(
+        array,
+        tg_dtype=dtypes.float32,
+        fn=lambda t: (stash.setdefault("u", t.reshape(2, 8)), (t + 2).sum())[1],
+      )
+
+  def test_scoped_handoff_rejects_returning_borrowed_tensor(self):
+    array = mx.array(np.arange(8, dtype=np.float32), dtype=mx.float32)
+    pool = MlxToTinygradLeasePool.from_mlx(array, tg_dtype=dtypes.float32, capacity=1)
+
+    with self.assertRaisesRegex(RuntimeError, "borrowed tensor directly"):
+      pool.run_with_mlx_tensor(array, lambda t: t)
+
+  def test_keyed_lease_pools_bucket_by_contract(self):
+    flat = mx.array(np.arange(16, dtype=np.float32), dtype=mx.float32)
+    matrix = mx.array(np.arange(16, dtype=np.float32).reshape(4, 4), dtype=mx.float32)
+    pools = MlxToTinygradLeasePools(capacity_per_key=2, synchronize_on_release=False)
+
+    with pools.acquire_from_mlx(flat, tg_dtype=dtypes.float32) as flat_lease:
+      self.assertEqual(flat_lease.key, MlxToTinygradLeaseKey((16,), "float32", dtypes.float32.base.name, 0))
+
+    with pools.acquire_from_mlx(matrix, tg_dtype=dtypes.float32) as matrix_lease:
+      self.assertEqual(matrix_lease.key, MlxToTinygradLeaseKey((4, 4), "float32", dtypes.float32.base.name, 0))
+
+    self.assertEqual(pools.pool_count, 2)
+
+  def test_alias_pool_registry_eviction_is_bounded(self):
+    pools = MlxToTinygradLeasePools(capacity_per_key=1, max_pools=2, synchronize_on_release=True)
+    arrays = [
+      mx.array(np.arange(4, dtype=np.float32), dtype=mx.float32),
+      mx.array(np.arange(6, dtype=np.float32).reshape(2, 3), dtype=mx.float32),
+      mx.array(np.arange(8, dtype=np.float32).reshape(2, 2, 2), dtype=mx.float32),
+    ]
+    keys = [MlxToTinygradLeaseKey.from_storage(mx.metal._unsafe_export_storage(arr), tg_dtype=dtypes.float32) for arr in arrays]
+
+    for arr in arrays:
+      result = pools.run_with_mlx_tensor(arr, tg_dtype=dtypes.float32, fn=lambda t: (t + 1).sum())
+      self.assertIsInstance(result, Tensor)
+
+    self.assertEqual(pools.pool_count, 2)
+    self.assertIsNone(pools.get_pool(keys[0]))
+    self.assertIsNotNone(pools.get_pool(keys[1]))
+    self.assertIsNotNone(pools.get_pool(keys[2]))
+
+  def test_copy_pool_scoped_handoff_realizes_returned_tensor(self):
+    first = mx.array(np.arange(16, dtype=np.float32), dtype=mx.float32)
+    second = mx.array(np.arange(16, dtype=np.float32) + np.float32(3), dtype=mx.float32)
+    pools = MlxToTinygradCopyLeasePools(capacity_per_key=1, synchronize_on_release=True)
+
+    result_first = pools.run_with_mlx_tensor(first, tg_dtype=dtypes.float32, fn=lambda t: (t + 1).sum())
+    result_second = pools.run_with_mlx_tensor(second, tg_dtype=dtypes.float32, fn=lambda t: (t + 1).sum())
+
+    np.testing.assert_allclose(result_first.numpy(), np.array((np.arange(16, dtype=np.float32) + 1).sum(), dtype=np.float32))
+    np.testing.assert_allclose(result_second.numpy(), np.array((np.arange(16, dtype=np.float32) + 4).sum(), dtype=np.float32))
+
+  def test_copy_pool_registry_eviction_is_bounded(self):
+    pools = MlxToTinygradCopyLeasePools(capacity_per_key=1, max_pools=2, synchronize_on_release=True)
+    arrays = [
+      mx.array(np.arange(4, dtype=np.float32), dtype=mx.float32),
+      mx.array(np.arange(6, dtype=np.float32).reshape(2, 3), dtype=mx.float32),
+      mx.array(np.arange(8, dtype=np.float32).reshape(2, 2, 2), dtype=mx.float32),
+    ]
+    keys = [MlxToTinygradCopyKey.from_storage(mx.metal._unsafe_export_storage(arr), tg_dtype=dtypes.float32) for arr in arrays]
+
+    for arr in arrays:
+      result = pools.run_with_mlx_tensor(arr, tg_dtype=dtypes.float32, fn=lambda t: (t + 1).sum())
+      self.assertIsInstance(result, Tensor)
+
+    self.assertEqual(pools.pool_count, 2)
+    self.assertIsNone(pools.get_pool(keys[0]))
+    self.assertIsNotNone(pools.get_pool(keys[1]))
+    self.assertIsNotNone(pools.get_pool(keys[2]))
+
+  def test_copy_pool_capacity_requires_release(self):
+    array = mx.array(np.arange(16, dtype=np.float32), dtype=mx.float32)
+    pool = MlxToTinygradCopyLeasePool.from_mlx(array, tg_dtype=dtypes.float32, capacity=1)
+
+    lease = pool.acquire_from_mlx(array)
+    with self.assertRaisesRegex(RuntimeError, "still in use"):
+      pool.acquire_from_mlx(array)
+    lease.release(synchronize=False)
+
+
+if __name__ == "__main__":
+  unittest.main()

mlx_tinygrad_interop/test_interop.py

@@ -0,0 +1,184 @@
+import unittest
+
+import mlx.core as mx
+import numpy as np
+from tinygrad import Tensor, dtypes
+
+from mlx_tinygrad_interop.stress_interop import apply_numpy_ops
+
+
+class TestStressHarnessNumerics(unittest.TestCase):
+  def test_numpy_matmul_lastdim_uses_reliable_baseline(self):
+    lhs = np.arange(2 * 2 * 4 * 31, dtype=np.float32).reshape(2, 2, 4, 31)
+    lhs = np.transpose(lhs, (3, 0, 2, 1)).reshape(2, 2, 4, 31)
+    weight = np.arange(31 * 7, dtype=np.float32).reshape(31, 7) / np.float32(17)
+
+    actual = apply_numpy_ops(lhs, [("matmul_lastdim", weight)])
+    expected = np.einsum("ik,kj->ij", lhs.reshape(-1, lhs.shape[-1]), weight, optimize=True)
+
+    np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-6)
+
+
+@unittest.skipUnless(mx.metal.is_available(), "Metal is not available")
+class TestMlxTinygradInterop(unittest.TestCase):
+  def test_mlx_slice_import_uses_backing_buffer_bytes(self):
+    backing = np.arange(4096, dtype=np.float32)
+    base = mx.array(backing, dtype=mx.float32)
+    view = base[16:1808]
+    mx.eval(view)
+
+    storage = mx.metal._unsafe_export_storage(view)
+    self.assertNotIn("nbytes", storage)
+    self.assertEqual(int(storage["offset_bytes"]), 64)
+    self.assertEqual(int(storage["logical_nbytes"]), 7168)
+    self.assertEqual(int(storage["buffer_nbytes"]), 16384)
+
+    tensor = Tensor._unsafe_from_metal_buffer_fast(
+      int(storage["mtl_buffer_ptr"]),
+      tuple(storage["shape"]),
+      dtype=dtypes.float32,
+      byte_offset=int(storage["offset_bytes"]),
+      buffer_nbytes=int(storage["buffer_nbytes"]),
+      owner=view,
+    )
+    np.testing.assert_array_equal(tensor.numpy(), backing[16:1808])
+
+  def test_single_entry_mlx_to_tinygrad_fast(self):
+    values = np.arange(128, dtype=np.float32)
+    array = mx.array(values, dtype=mx.float32)
+    tensor = mx.metal._unsafe_to_tinygrad_fast(array, dtypes.float32, owner=array)
+    np.testing.assert_array_equal(tensor.numpy(), values)
+
+  def test_fast_import_handles_zero_offset_oversized_backing_buffer(self):
+    backing = np.arange(2048, dtype=np.float32)
+    view = mx.array(backing, dtype=mx.float32)[:1792].reshape(7, 256)
+    mx.eval(view)
+
+    storage = mx.metal._unsafe_export_storage(view)
+    self.assertEqual(int(storage["offset_bytes"]), 0)
+    self.assertGreater(int(storage["buffer_nbytes"]), int(storage["logical_nbytes"]))
+
+    tensor = Tensor._unsafe_from_metal_buffer_fast(
+      int(storage["mtl_buffer_ptr"]),
+      tuple(storage["shape"]),
+      dtype=dtypes.float32,
+      byte_offset=int(storage["offset_bytes"]),
+      buffer_nbytes=int(storage["buffer_nbytes"]),
+      owner=view,
+    )
+    np.testing.assert_array_equal(tensor.numpy(), backing[:1792].reshape(7, 256))
+
+  def test_rebindable_slot_mutates_previous_reference(self):
+    first = np.arange(64, dtype=np.float32)
+    second = first + np.float32(1)
+    array_a = mx.array(first, dtype=mx.float32)
+    array_b = mx.array(second, dtype=mx.float32)
+    storage_a = mx.metal._unsafe_export_storage(array_a)
+    borrower = Tensor._unsafe_metal_borrower(
+      int(storage_a["mtl_buffer_ptr"]),
+      tuple(storage_a["shape"]),
+      dtype=dtypes.float32,
+      byte_offset=int(storage_a["offset_bytes"]),
+      buffer_nbytes=int(storage_a["buffer_nbytes"]),
+      owner=array_a,
+    )
+
+    tensor_a = mx.metal._unsafe_rebind_tinygrad(array_a, borrower, owner=array_a)
+    tensor_b = mx.metal._unsafe_rebind_tinygrad(array_b, borrower, owner=array_b)
+
+    self.assertIs(tensor_a, tensor_b)
+    np.testing.assert_array_equal(tensor_a.numpy(), second)
+    np.testing.assert_array_equal(tensor_b.numpy(), second)
+
+  def test_rebindable_slot_updates_external_ptr_metadata(self):
+    first = np.arange(32, dtype=np.float32)
+    second = first + np.float32(2)
+    array_a = mx.array(first, dtype=mx.float32)
+    array_b = mx.array(second, dtype=mx.float32)
+    storage_a = mx.metal._unsafe_export_storage(array_a)
+    storage_b = mx.metal._unsafe_export_storage(array_b)
+    borrower = Tensor._unsafe_metal_borrower(
+      int(storage_a["mtl_buffer_ptr"]),
+      tuple(storage_a["shape"]),
+      dtype=dtypes.float32,
+      byte_offset=int(storage_a["offset_bytes"]),
+      buffer_nbytes=int(storage_a["buffer_nbytes"]),
+      owner=array_a,
+    )
+
+    mx.metal._unsafe_rebind_tinygrad(array_b, borrower, owner=array_b)
+
+    self.assertIsNotNone(borrower._base_buf.options)
+    self.assertEqual(borrower._base_buf.options.external_ptr, int(storage_b["mtl_buffer_ptr"]))
+
+  def test_two_slots_hold_two_distinct_snapshots_until_reused(self):
+    first = np.arange(16, dtype=np.float32)
+    second = first + np.float32(1)
+    third = first + np.float32(2)
+    array_a = mx.array(first, dtype=mx.float32)
+    array_b = mx.array(second, dtype=mx.float32)
+    array_c = mx.array(third, dtype=mx.float32)
+    storage_a = mx.metal._unsafe_export_storage(array_a)
+    storage_b = mx.metal._unsafe_export_storage(array_b)
+    slot_a = Tensor._unsafe_metal_borrower(
+      int(storage_a["mtl_buffer_ptr"]),
+      tuple(storage_a["shape"]),
+      dtype=dtypes.float32,
+      byte_offset=int(storage_a["offset_bytes"]),
+      buffer_nbytes=int(storage_a["buffer_nbytes"]),
+      owner=array_a,
+    )
+    slot_b = Tensor._unsafe_metal_borrower(
+      int(storage_b["mtl_buffer_ptr"]),
+      tuple(storage_b["shape"]),
+      dtype=dtypes.float32,
+      byte_offset=int(storage_b["offset_bytes"]),
+      buffer_nbytes=int(storage_b["buffer_nbytes"]),
+      owner=array_b,
+    )
+
+    tensor_a = mx.metal._unsafe_rebind_tinygrad(array_a, slot_a, owner=array_a)
+    tensor_b = mx.metal._unsafe_rebind_tinygrad(array_b, slot_b, owner=array_b)
+    mx.metal._unsafe_rebind_tinygrad(array_c, slot_a, owner=array_c)
+
+    self.assertIsNot(tensor_a, tensor_b)
+    np.testing.assert_array_equal(tensor_a.numpy(), third)
+    np.testing.assert_array_equal(tensor_b.numpy(), second)
+
+  def test_rebindable_slot_rejects_shape_mismatch(self):
+    base = np.arange(32, dtype=np.float32)
+    array_ok = mx.array(base, dtype=mx.float32)
+    array_bad = mx.array(base.reshape(8, 4), dtype=mx.float32)
+    storage = mx.metal._unsafe_export_storage(array_ok)
+    borrower = Tensor._unsafe_metal_borrower(
+      int(storage["mtl_buffer_ptr"]),
+      tuple(storage["shape"]),
+      dtype=dtypes.float32,
+      byte_offset=int(storage["offset_bytes"]),
+      buffer_nbytes=int(storage["buffer_nbytes"]),
+      owner=array_ok,
+    )
+
+    with self.assertRaisesRegex(ValueError, "shape="):
+      mx.metal._unsafe_rebind_tinygrad(array_bad, borrower, owner=array_bad)
+
+  def test_rebindable_slot_rejects_dtype_mismatch(self):
+    base = np.arange(32, dtype=np.float32)
+    array_ok = mx.array(base, dtype=mx.float32)
+    array_bad = mx.array(base.view(np.int32), dtype=mx.int32)
+    storage = mx.metal._unsafe_export_storage(array_ok)
+    borrower = Tensor._unsafe_metal_borrower(
+      int(storage["mtl_buffer_ptr"]),
+      tuple(storage["shape"]),
+      dtype=dtypes.float32,
+      byte_offset=int(storage["offset_bytes"]),
+      buffer_nbytes=int(storage["buffer_nbytes"]),
+      owner=array_ok,
+    )
+
+    with self.assertRaisesRegex(ValueError, "dtype="):
+      mx.metal._unsafe_rebind_tinygrad(array_bad, borrower, owner=array_bad)
+
+
+if __name__ == "__main__":
+  unittest.main()

nix/darwin-build-fixes.patch

@@ -59,7 +59,7 @@ index 262b0495..5c7446ad 100644
 
  add_custom_command(
    OUTPUT ${MLX_METAL_PATH}/mlx.metallib
--  COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o
+-  COMMAND xcrun -sdk macosx metal ${KERNEL_AIR} -o
 +  COMMAND metallib ${KERNEL_AIR} -o
            ${MLX_METAL_PATH}/mlx.metallib
    DEPENDS ${KERNEL_AIR}

nix/mlx.nix

@@ -1,158 +0,0 @@
-{ stdenv
-, lib
-, fetchFromGitHub
-, replaceVars
-, fetchzip
-, cmake
-, nlohmann_json
-, apple-sdk_26
-, metal-toolchain
-, runCommand
-, fmt
-, python313Packages
-, uvLockMlxVersion
-, uvLockMlxRev
-}:
-
-assert stdenv.isDarwin;
-
-let
-  python = python313Packages.python;
-
-  # Static dependencies included directly during compilation
-  gguf-tools = fetchFromGitHub {
-    owner = "antirez";
-    repo = "gguf-tools";
-    rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
-    hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
-  };
-
-  metal_cpp = fetchzip {
-    url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
-    hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
-  };
-
-  nanobind = fetchFromGitHub {
-    owner = "wjakob";
-    repo = "nanobind";
-    rev = "v2.10.2";
-    hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
-    fetchSubmodules = true;
-  };
-
-  mlx = stdenv.mkDerivation rec {
-    pname = "mlx";
-    version = uvLockMlxVersion;
-    pyproject = true;
-
-    src = fetchFromGitHub {
-      owner = "rltakashige";
-      repo = "mlx-jaccl-fix-small-recv";
-      rev = uvLockMlxRev;
-      hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
-    };
-
-    patches = [
-      (replaceVars ./darwin-build-fixes.patch {
-        sdkVersion = apple-sdk_26.version;
-        metalVersion = metal-toolchain.metalVersion;
-      })
-    ];
-
-    postPatch = ''
-      substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
-        --replace-fail "g++" "$CXX"
-    '';
-
-    dontUseCmakeConfigure = true;
-
-    enableParallelBuilding = true;
-
-    # Allows multiple cores to be used in Python builds.
-    postUnpack = ''
-      export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}"
-    '';
-
-    # Updates the wrong fetcher rev attribute
-    passthru.skipBulkUpdate = true;
-
-    env = {
-      DEV_RELEASE = 1;
-      CMAKE_ARGS = toString [
-        (lib.cmakeBool "USE_SYSTEM_FMT" true)
-        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
-        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}")
-        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
-        (lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
-        (lib.cmakeBool "MLX_BUILD_CPU" true)
-        (lib.cmakeBool "MLX_BUILD_METAL" true)
-        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
-        (lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${apple-sdk_26.version}")
-        (lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${apple-sdk_26.passthru.sdkroot}")
-      ];
-      SDKROOT = apple-sdk_26.passthru.sdkroot;
-      MACOSX_DEPLOYMENT_TARGET = apple-sdk_26.version;
-    };
-
-    build-system = [
-      python313Packages.setuptools
-    ];
-
-    nativeBuildInputs = [
-      cmake
-      metal-toolchain
-      python313Packages.pypaBuildHook
-      python313Packages.pypaInstallHook
-      python313Packages.setuptools
-      python313Packages.typing-extensions
-      python313Packages.wheel
-      python313Packages.cmake
-      python313Packages.ninja
-    ];
-
-    buildInputs = [
-      fmt
-      gguf-tools
-      python313Packages.nanobind
-      python313Packages.pybind11
-      apple-sdk_26
-    ];
-
-    # Tests require Metal GPU access which isn't available in the Nix sandbox.
-    # To run tests, build with: nix build --option sandbox false .#mlx.passthru.tests.mlxTest
-    doCheck = false;
-
-    pythonImportsCheck = [ "mlx" ];
-
-    passthru.tests = {
-      # Runs example scripts to verify MLX works. Requires --option sandbox false
-      # since Metal GPU access is needed.
-      mlxTest =
-        runCommand "run-mlx-examples"
-          {
-            buildInputs = [ mlx ];
-            nativeBuildInputs = [ python ];
-          }
-          ''
-            cp ${src}/examples/python/logistic_regression.py .
-            ${python.interpreter} logistic_regression.py
-            rm logistic_regression.py
-
-            cp ${src}/examples/python/linear_regression.py .
-            ${python.interpreter} linear_regression.py
-            rm linear_regression.py
-
-            touch $out
-          '';
-    };
-
-    meta = {
-      homepage = "https://github.com/ml-explore/mlx";
-      description = "Array framework for Apple silicon";
-      changelog = "https://github.com/ml-explore/mlx/releases/tag/${src.tag}";
-      license = lib.licenses.mit;
-      platforms = [ "aarch64-darwin" ];
-    };
-  };
-in
-mlx

packaging/pyinstaller/exo.spec

@@ -1,5 +1,6 @@
 # -*- mode: python ; coding: utf-8 -*-
 
+import sys
 import importlib.util
 import shutil
 from pathlib import Path
@@ -56,6 +57,7 @@ HIDDEN_IMPORTS = sorted(
     set(
         collect_submodules("mlx")
         + _safe_collect("mlx_lm")
+        + _safe_collect("mlx_vlm")
         + _safe_collect("transformers")
     )
 )
@@ -67,16 +69,19 @@ DATAS: list[tuple[str, str]] = [
     (str(EXO_SHARED_MODELS_DIR), "exo/shared/models"),
 ]
 
-MACMON_PATH = shutil.which("macmon")
-if MACMON_PATH is None:
-    raise SystemExit(
-        "macmon binary not found in PATH. "
-        "Install it via: brew install macmon"
-    )
+if sys.platform == "darwin":
+    MACMON_PATH = shutil.which("macmon")
+    if MACMON_PATH is None:
+        raise SystemExit(
+            "macmon binary not found in PATH. "
+            "Install the pinned fork used by exo via: "
+            "cargo install --git https://github.com/vladkens/macmon "
+            "--rev a1cd06b6cc0d5e61db24fd8832e74cd992097a7d macmon --force"
+        )
 
 BINARIES: list[tuple[str, str]] = [
     (MACMON_PATH, "."),
-]
+] if sys.platform == "darwin" else []
 
 a = Analysis(
     [str(ENTRYPOINT)],
@@ -120,4 +125,3 @@ coll = COLLECT(
     upx_exclude=[],
     name="exo",
 )
-

pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 name = "exo"
-version = "0.3.68"
+version = "0.3.70"
 description = "Exo"
 readme = "README.md"
-requires-python = ">=3.13"
+requires-python = "==3.13.*"
 dependencies = [
   "aiofiles>=24.1.0",
   "aiohttp>=3.12.14",
@@ -12,23 +12,26 @@ dependencies = [
   "fastapi>=0.116.1",
   "filelock>=3.18.0",
   "rustworkx>=0.17.1",
-  "huggingface-hub>=0.33.4",
+  "huggingface-hub>=1.8.0",
   "psutil>=7.0.0",
   "loguru>=0.7.3",
-  "exo_pyo3_bindings",                         # rust bindings
+  "exo-pyo3-bindings",                       # rust bindings
   "anyio==4.11.0",
-  "mlx; sys_platform == 'darwin'",
-  "mlx[cpu]==0.30.6; sys_platform == 'linux'",
-  "mlx-lm",
-  "tiktoken>=0.12.0",                          # required for kimi k2 tokenizer
+  "mlx==0.31.2; sys_platform == 'darwin'",
+  "mlx-lm; sys_platform=='darwin'",
+  "tiktoken>=0.12.0",                        # required for kimi k2 tokenizer
   "hypercorn>=0.18.0",
   "openai-harmony>=0.0.8",
   "httpx>=0.28.1",
   "tomlkit>=0.14.0",
-  "mflux==0.16.9",
+  "mflux==0.17.2; sys_platform == 'darwin'",
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
+  "mlx-vlm>=0.3.11",
+  "transformers>=5.0.0,<5.4.0",
+  "tinygrad",                                # TODO: maybe add a version or something idk lol
+  "torch>=2.10.0",
 ]
 
 [project.scripts]
@@ -45,11 +48,26 @@ dev = [
   "ruff>=0.11.13",
 ]
 
-# mlx[cuda] requires a newer version of mlx. the ideal on linux is: default to mlx[cpu] unless[cuda] specified.
 [project.optional-dependencies]
-# cuda = [
-#     "mlx[cuda]==0.26.3",
-# ]
+build = ["nanobind"]
+cpu = [
+  "mlx==0.31.1; sys_platform == 'linux'",
+  "mlx-cpu==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
+  "torch>=2.10.0; sys_platform == 'linux'",
+]
+cuda12 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
+  "mlx-cuda-12==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
+  "torch>=2.10.0; sys_platform == 'linux'",
+]
+cuda13 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
+  "mlx-cuda-13==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
+  "torch>=2.10.0; sys_platform == 'linux'",
+]
 
 ###
 # workspace configuration
@@ -59,12 +77,31 @@ dev = [
 members = ["rust/exo_pyo3_bindings", "bench"]
 
 [tool.uv.sources]
-exo_pyo3_bindings = { workspace = true }
-mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
-# Uncomment to use local mlx/mlx-lm development versions:
-# mlx = { path = "/Users/Shared/mlx", editable=true }
-# mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
+exo-pyo3-bindings = { workspace = true }
+mlx = { git = "https://github.com/AndreiCravtov/mlx.git", branch = "andrei/mlx-tinygrad-convert", marker = "sys_platform == 'darwin'" }
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arrayscache-leak" }
+torch = [
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
+  { index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },
+  { index = "pytorch-cpu", marker = "(extra != 'cuda12' and extra != 'cuda13' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+]
+vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
+tinygrad = { git = "https://github.com/AndreiCravtov/tinygrad.git", branch = "andrei/tinygrad-mlx-convert" }
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu120"
+url = "https://download.pytorch.org/whl/cu120"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
 
 [build-system]
 requires = ["uv_build>=0.8.9,<0.9.0"]
@@ -75,7 +112,7 @@ build-backend = "uv_build"
 ###
 
 [tool.basedpyright]
-include = [".venv/lib/mlx", ".venv/lib/mlx_lm", "src", "bench"]
+include = ["src", "bench"]
 typeCheckingMode = "strict"
 failOnWarnings = true
 
@@ -102,9 +139,14 @@ exclude = [
 ]
 stubPath = ".mlx_typings"
 
+[[tool.basedpyright.executionEnvironments]]
+root = "src/exo/worker/engines/image"
+reportMissingModuleSource = false
+
 [[tool.basedpyright.executionEnvironments]]
 root = "src"
 
+
 ###
 # uv configuration
 ###
@@ -114,6 +156,51 @@ root = "src"
 required-version = ">=0.8.6"
 prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
+conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }, { extra = "cpu" }]]
+constraint-dependencies = ["transformers>=5.0.0,<5.4.0"]
+override-dependencies = [
+  "mlx==0.31.1; sys_platform=='linux'",
+  "mlx; sys_platform=='darwin'",
+]
+
+[tool.uv.extra-build-dependencies]
+miniaudio = ["setuptools", "cffi", "pycparser"]
+mlx = [
+  "setuptools",
+  "typing-extensions",
+  "nanobind",
+  "pybind11",
+  "wheel",
+  "cmake",
+  "ninja",
+]
+mlx-lm = ["setuptools"]
+tinygrad = ["setuptools"]
+xgrammar = [
+  "nanobind",
+  "setuptools",
+  "scikit-build-core",
+  "packaging",
+  "pathspec",
+]
+rouge-score = ["setuptools"]
+sacrebleu = ["setuptools"]
+sqlitedict = ["setuptools"]
+word2number = ["setuptools"]
+vllm = [
+  "setuptools",
+  "setuptools-scm",
+  "scikit-build-core",
+  "jinja2",
+  "wheel",
+  "markupsafe",
+  "typing-extensions",
+  "torch",
+]
+fastsafetensors = ["setuptools", "pybind11"]
+torch = ["typing-extensions"]
+torchvision = ["torch"]
+torchaudio = ["torch"]
 
 ###
 # ruff configuration
@@ -121,7 +208,6 @@ environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 
 [tool.ruff]
 extend-exclude = [
-  "shared/protobufs/**",
   "*mlx_typings/**",
   "rust/exo_pyo3_bindings/**",
   "bench/vendor/**",

python/parts.nix

@@ -1,18 +1,36 @@
 { inputs, ... }:
-{
-  perSystem =
-    { config, self', pkgs, lib, system, ... }:
+let
+  # Load workspace from uv.lock
+  workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
+    workspaceRoot = ../.;
+  };
+
+  mkPythonSet = { pkgs, lib, self', members }:
     let
-      # Load workspace from uv.lock
-      workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
-        workspaceRoot = inputs.self;
-      };
-
-      # Create overlay from workspace
-      # Use wheels from PyPI for most packages; we override mlx with our pure Nix Metal build
-      overlay = workspace.mkPyprojectOverlay { sourcePreference = "wheel"; };
-
-      # Override overlay to inject Nix-built components
+      inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
+      inherit (pkgs.config) cudaSupport;
+      inherit (pkgs) cudaPackages;
+      cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
+      libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
+      python = pkgs.python313;
+      cudaLibs = with cudaPackages; [
+        cuda_cudart
+        cuda_cccl
+        cuda_cupti
+        cuda_nvrtc
+        cuda_nvtx
+        cudnn
+        libcufile
+        libcublas
+        libcufft
+        libcurand
+        libcusolver
+        libcusparse
+        libcusparse_lt
+        libnvjitlink
+        libnvshmem
+        nccl
+      ];
       exoOverlay = final: prev: {
         # Replace workspace exo_pyo3_bindings with Nix-built wheel.
         # Preserve passthru so mkVirtualEnv can resolve dependency groups.
@@ -32,126 +50,162 @@
           '';
         };
       };
+      buildSystemsOverlay = final: prev:
+        lib.optionalAttrs isDarwin
+          {
+            mlx = prev.mlx.overrideAttrs (old:
+              let
+                # Static dependencies included directly during compilation
+                gguf-tools = pkgs.fetchFromGitHub {
+                  owner = "antirez";
+                  repo = "gguf-tools";
+                  rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
+                  hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
+                };
 
-      python = pkgs.python313;
+                metal_cpp = pkgs.fetchzip {
+                  url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
+                  hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
+                };
 
-      # Overlay to provide build systems and custom packages
-      buildSystemsOverlay = final: prev: {
-        # mlx-lm is a git dependency that needs setuptools
-        mlx-lm = prev.mlx-lm.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        # rouge-score and sacrebleu don't declare setuptools as a build dependency
-        rouge-score = prev.rouge-score.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        sacrebleu = prev.sacrebleu.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        sqlitedict = prev.sqlitedict.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        word2number = prev.word2number.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-      } // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
-        # Use our pure Nix-built MLX with Metal support (macOS only)
-        mlx = self'.packages.mlx;
+                nanobind = pkgs.fetchFromGitHub {
+                  owner = "wjakob";
+                  repo = "nanobind";
+                  rev = "v2.10.2";
+                  hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
+                  fetchSubmodules = true;
+                };
+              in
+              {
+                nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake self'.packages.metal-toolchain ];
+                # TODO: non-sdk_26 support
+                buildInputs = (old.buildInputs or [ ])
+                ++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
+                patches = [
+                  (pkgs.replaceVars ../nix/darwin-build-fixes.patch {
+                    sdkVersion = pkgs.apple-sdk_26.version;
+                    inherit (self'.packages.metal-toolchain) metalVersion;
+                  })
+                ];
+                postPatch = ''
+                  substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
+                    --replace-fail "g++" "${lib.getExe' pkgs.stdenv.cc "c++"}"
+                '';
+
+                DEV_RELEASE = 1;
+                CMAKE_ARGS = toString ([
+                  (lib.cmakeBool "USE_SYSTEM_FMT" true)
+                  (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
+                  (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${pkgs.nlohmann_json.src}")
+                  (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
+                  (lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
+                  (lib.cmakeBool "MLX_BUILD_CPU" true)
+                  (lib.cmakeBool "MLX_BUILD_METAL" true)
+                  (lib.cmakeOptionType "string" "CMAKE_INSTALL_LIBDIR" "lib")
+                  (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
+                  (lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${pkgs.apple-sdk_26.version}")
+                  (lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${pkgs.apple-sdk_26.passthru.sdkroot}")
+                ] ++ lib.optionals (isDarwin && isx86_64) [
+                  (lib.cmakeBool "MLX_ENABLE_X64_MAC" true)
+                ]);
+                SDKROOT = pkgs.apple-sdk_26.passthru.sdkroot;
+                MACOSX_DEPLOYMENT_TARGET = pkgs.apple-sdk_26.version;
+              });
+          } // lib.optionalAttrs isLinux {
+          mlx = prev.mlx.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
+            autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
+            postInstall = ''
+              cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
+            '';
+          });
+        } // lib.optionalAttrs cudaSupport {
+          "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ cudaLibs ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          torch = prev.torch.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+        };
+      pyprojectOverlay = workspace.mkPyprojectOverlay {
+        sourcePreference = "wheel";
+        dependencies = members;
+      };
+      editableOverlay = workspace.mkEditablePyprojectOverlay {
+        # Use environment variable pointing to editable root directory
+        root = "$REPO_ROOT";
+        members = [ "exo" "exo-bench" ];
       };
-
-      # Additional overlay for Linux-specific fixes (type checking env).
-      # Native wheels have shared lib dependencies we don't need at type-check time.
-      linuxOverlay = final: prev:
-        let
-          ignoreMissing = drv: drv.overrideAttrs { autoPatchelfIgnoreMissingDeps = [ "*" ]; };
-          nvidiaPackages = lib.filterAttrs (name: _: lib.hasPrefix "nvidia-" name) prev;
-        in
-        lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux (
-          (lib.mapAttrs (_: ignoreMissing) nvidiaPackages) // {
-            mlx = ignoreMissing prev.mlx;
-            mlx-cuda-13 = prev.mlx-cuda-13.overrideAttrs (old: {
-              buildInputs = (old.buildInputs or [ ]) ++ [
-                final.nvidia-cublas
-                final.nvidia-cuda-nvrtc
-                final.nvidia-cudnn-cu13
-                final.nvidia-nccl-cu13
-              ];
-              preFixup = ''
-                addAutoPatchelfSearchPath ${final.nvidia-cublas}
-                addAutoPatchelfSearchPath ${final.nvidia-cuda-nvrtc}
-                addAutoPatchelfSearchPath ${final.nvidia-cudnn-cu13}
-                addAutoPatchelfSearchPath ${final.nvidia-nccl-cu13}
-              '';
-              autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-            });
-            torch = ignoreMissing prev.torch;
-            triton = ignoreMissing prev.triton;
-          }
-        );
-
       pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
         inherit python;
       }).overrideScope (
         lib.composeManyExtensions [
           inputs.pyproject-build-systems.overlays.default
-          overlay
+          pyprojectOverlay
           exoOverlay
           buildSystemsOverlay
-          linuxOverlay
         ]
       );
-      # mlx-cpu and mlx-cuda-13 both ship mlx/ site-packages files; keep first.
-      # mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13 ship overlapping files.
-      venvCollisionPaths = lib.optionals pkgs.stdenv.hostPlatform.isLinux [
-        "lib/python3.13/site-packages/mlx*"
-        "lib/python3.13/site-packages/nvidia*"
-      ];
-
-      # Exclude bench deps from main env (bench has its own benchVenv)
-      exoDeps = removeAttrs workspace.deps.default [ "exo-bench" ];
-
-      exoVenv = (pythonSet.mkVirtualEnv "exo-env" exoDeps).overrideAttrs {
-        venvIgnoreCollisions = venvCollisionPaths;
-      };
-
-      # Virtual environment with dev dependencies for testing
-      testVenv = (pythonSet.mkVirtualEnv "exo-test-env" (
-        exoDeps // {
-          exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
-        }
-      )).overrideAttrs {
-        venvIgnoreCollisions = venvCollisionPaths;
-      };
-
-      mkPythonScript = name: path: pkgs.writeShellApplication {
+      venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
+      mkApp = cmd: name: pkgs.writeShellApplication {
         inherit name;
-        runtimeInputs = [ exoVenv ];
         runtimeEnv = {
           EXO_DASHBOARD_DIR = self'.packages.dashboard;
           EXO_RESOURCES_DIR = inputs.self + /resources;
         };
-        text = ''exec python ${path} "$@"'';
+        runtimeInputs = [
+          # mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
+          (venv name)
+        ]
+        ++ lib.optionals isDarwin [ pkgs.macmon ];
+        text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
       };
+    in
+    {
+      inherit venv;
+      editablePythonSet = pythonSet.overrideScope editableOverlay;
+      mkPythonScript = path: mkApp ''python ${path} "$@"'';
+      mkExo = mkApp ''exo "$@"'';
+    };
+in
+{
+  perSystem =
+    { self', pkgs, unfreePkgs, lib, ... }:
+    let
+      inherit (pkgs.stdenv.hostPlatform) isLinux;
+      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "cpu" ]; }; }) editablePythonSet mkExo;
 
-      benchVenv = pythonSet.mkVirtualEnv "exo-bench-env" {
-        exo-bench = [ ];
+      # Virtual environment with dev dependencies for testing
+      testVenv = (mkPythonSet {
+        inherit self' pkgs lib; members = {
+        exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
       };
+      }).venv "exo-test";
 
-      mkBenchScript = name: path: pkgs.writeShellApplication {
-        inherit name;
-        runtimeInputs = [ benchVenv ];
-        text = ''exec python ${path} "$@"'';
+      mkBenchScript = (mkPythonSet {
+        inherit self' pkgs lib; members = {
+        exo = [ "cpu" ];
+        exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
       };
+      }).mkPythonScript;
 
       mkSimplePythonScript = name: path: pkgs.writeShellApplication {
         inherit name;
@@ -159,56 +213,35 @@
         text = ''exec python ${path} "$@"'';
       };
 
-      exoPackage = pkgs.runCommand "exo"
-        {
-          nativeBuildInputs = [ pkgs.makeWrapper ];
-        }
-        ''
-          mkdir -p $out/bin
-
-          # Create wrapper script
-          makeWrapper ${exoVenv}/bin/exo $out/bin/exo \
-            --set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
-            --set EXO_RESOURCES_DIR ${inputs.self + /resources} \
-            ${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
-        '';
     in
     {
-      # Python package only available on macOS (requires MLX/Metal)
-      packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
-        {
-          exo = exoPackage;
-          # Test environment for running pytest outside of Nix sandbox (needs GPU access)
-          exo-test-env = testVenv;
-        } // {
+      packages = {
+        exo = mkExo "exo";
+        editableVenv = editablePythonSet.mkVirtualEnv "exo-dev-env" { exo = [ "dev" ]; };
+        # for running tests in ci
+        exo-test-env = testVenv;
         exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
         exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
         exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
+        # used by ./tests/run_exo_on.sh
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
+      } // lib.optionalAttrs isLinux {
+        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "cuda12" ]; }; }).mkExo "exo-cuda-12";
+        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "cuda13" ]; }; }).mkExo "exo-cuda-13";
       };
 
       checks = {
-        # Ruff linting (works on all platforms)
         lint = pkgs.runCommand "ruff-lint" { } ''
           export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
           ${pkgs.ruff}/bin/ruff check ${inputs.self}
           touch $out
         '';
 
-        # Hermetic basedpyright type checking
-        typecheck = pkgs.runCommand "typecheck"
-          {
-            nativeBuildInputs = [
-              testVenv
-              pkgs.basedpyright
-            ];
-          }
-          ''
-            cd ${inputs.self}
-            export HOME=$TMPDIR
-            basedpyright --pythonpath ${testVenv}/bin/python
-            touch $out
-          '';
+        typecheck = pkgs.runCommand "typecheck" { nativeBuildInputs = [ testVenv ]; } ''
+          cd ${inputs.self}
+          basedpyright
+          touch $out
+        '';
       };
     };
 }

resources/inference_model_cards/mlx-community--DeepSeek-V3.1-4bit.toml

@@ -9,5 +9,13 @@ quantization = "4bit"
 base_model = "DeepSeek V3.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
 
+context_length = 131072
+
 [storage_size]
 in_bytes = 405874409472
+
+# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.1/blob/main/generation_config.json
+# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.1/discussions/19
+[sampling_defaults]
+temperature = 0.6
+top_p = 0.95

resources/inference_model_cards/mlx-community--DeepSeek-V3.1-8bit.toml

@@ -9,5 +9,13 @@ quantization = "8bit"
 base_model = "DeepSeek V3.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
 
+context_length = 131072
+
 [storage_size]
 in_bytes = 765577920512
+
+# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.1/blob/main/generation_config.json
+# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.1/discussions/19
+[sampling_defaults]
+temperature = 0.6
+top_p = 0.95

resources/inference_model_cards/mlx-community--DeepSeek-V3.2-4bit.toml

@@ -0,0 +1,21 @@
+model_id = "mlx-community/DeepSeek-V3.2-4bit"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 128
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "deepseek"
+quantization = "4bit"
+base_model = "DeepSeek V3.2"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+context_length = 131072
+
+[storage_size]
+in_bytes = 378086226621
+
+# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.2/blob/main/generation_config.json
+# Source: https://docs.vllm.ai/projects/recipes/en/latest/DeepSeek/DeepSeek-V3_2.html
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95

resources/inference_model_cards/mlx-community--DeepSeek-V3.2-8bit.toml

@@ -0,0 +1,21 @@
+model_id = "mlx-community/DeepSeek-V3.2-8bit"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 128
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "deepseek"
+quantization = "8bit"
+base_model = "DeepSeek V3.2"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+context_length = 131072
+
+[storage_size]
+in_bytes = 755957120916
+
+# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.2/blob/main/generation_config.json
+# Source: https://docs.vllm.ai/projects/recipes/en/latest/DeepSeek/DeepSeek-V3_2.html
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95

resources/inference_model_cards/mlx-community--GLM-4.5-Air-8bit.toml

@@ -9,5 +9,12 @@ quantization = "8bit"
 base_model = "GLM 4.5 Air"
 capabilities = ["text", "thinking", "thinking_toggle"]
 
+context_length = 131072
+
 [storage_size]
 in_bytes = 122406567936
+
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 0.6
+top_p = 0.95

resources/inference_model_cards/mlx-community--GLM-4.5-Air-bf16.toml

@@ -9,5 +9,12 @@ quantization = "bf16"
 base_model = "GLM 4.5 Air"
 capabilities = ["text", "thinking", "thinking_toggle"]
 
+context_length = 131072
+
 [storage_size]
 in_bytes = 229780750336
+
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 0.6
+top_p = 0.95

resources/inference_model_cards/mlx-community--GLM-4.7-4bit.toml

@@ -9,5 +9,14 @@ quantization = "4bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
 
+context_length = 202752
+
 [storage_size]
 in_bytes = 198556925568
+
+# Source: https://huggingface.co/zai-org/GLM-4.7
+# Source: https://unsloth.ai/docs/models/glm-4.7-flash
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95

resources/inference_model_cards/mlx-community--GLM-4.7-6bit.toml

@@ -9,5 +9,14 @@ quantization = "6bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
 
+context_length = 202752
+
 [storage_size]
 in_bytes = 286737579648
+
+# Source: https://huggingface.co/zai-org/GLM-4.7
+# Source: https://unsloth.ai/docs/models/glm-4.7-flash
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95

resources/inference_model_cards/mlx-community--GLM-4.7-8bit-gs32.toml

@@ -9,5 +9,14 @@ quantization = "8bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
 
+context_length = 202752
+
 [storage_size]
 in_bytes = 396963397248
+
+# Source: https://huggingface.co/zai-org/GLM-4.7
+# Source: https://unsloth.ai/docs/models/glm-4.7-flash
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-4bit.toml

@@ -9,5 +9,14 @@ quantization = "4bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
 
+context_length = 202752
+
 [storage_size]
 in_bytes = 19327352832
+
+# Source: https://huggingface.co/zai-org/GLM-4.7-Flash
+# Source: https://unsloth.ai/docs/models/glm-4.7-flash
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-5bit.toml

@@ -9,5 +9,14 @@ quantization = "5bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
 
+context_length = 202752
+
 [storage_size]
 in_bytes = 22548578304
+
+# Source: https://huggingface.co/zai-org/GLM-4.7-Flash
+# Source: https://unsloth.ai/docs/models/glm-4.7-flash
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95