pytorch · kiymetakdemir · Aug 7, 2026 · Aug 11, 2026 · Aug 11, 2026
diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml
@@ -678,3 +678,113 @@ jobs:
           exit 1
         fi
         echo "::endgroup::"
+
+  # Off-graph KV cache: the cache is a runtime object, so this path can only be
+  # exercised by the C++ runner (pybindings cannot bind a cache_key). Also the
+  # only coverage that the layout the export publishes is actually consumable.
+  test-mlx-llm-offgraph:
+    # Requires HuggingFace secrets — skip on fork PRs.
+    needs: run-decision
+    if: |
+      (github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request') &&
+      (github.event_name == 'pull_request' || needs.run-decision.outputs.is-full-run == 'true')
+    strategy:
+      fail-fast: false
+      matrix:
+        model:
+          - id: "unsloth/Llama-3.2-1B-Instruct"
+            name: "llama-1b"
+            chat: "llama3"
+            runner: "macos-14-xlarge"
+          - id: "unsloth/gemma-3-1b-it"
+            name: "gemma3-1b"
+            chat: "gemma"
+            runner: "macos-14-xlarge"
+          # Only model with KV sharing: 15 caches for 35 layers, and sliding and
+          # full-attention caches with different head shapes.
+          - id: "google/gemma-4-E2B-it"
+            name: "gemma4-e2b"
+            chat: "gemma4"
+            runner: "macos-15-xlarge"
+    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
+    secrets: inherit
+    with:
+      default-packages: ""
+      job-name: test-mlx-llm-offgraph-${{ matrix.model.name }}
+      runner: ${{ matrix.model.runner }}
+      python-version: "3.12"
+      submodules: recursive
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      secrets-env: EXECUTORCH_HF_TOKEN
+      # Higher than the sibling LLM job: this one also builds and installs
+      # ExecuTorch so the standalone runner project can link against it.
+      timeout: 120
+      script: |
+        set -eux
+        export HF_HUB_DISABLE_XET=1
+
+        MODEL_ID="${{ matrix.model.id }}"
+        MODEL_NAME="${{ matrix.model.name }}"
+        CHAT="${{ matrix.model.chat }}"
+
+        echo "::group::Install ExecuTorch and build the MLX runtime"
+        ${CONDA_RUN} python install_executorch.py > /dev/null
+        ${CONDA_RUN} cmake --preset mlx-release
+        ${CONDA_RUN} cmake --build cmake-out --target install -j$(( $(sysctl -n hw.ncpu) - 1 ))
+        echo "::endgroup::"
+
+        echo "::group::Build the C++ runner"
+        # Standalone find_package(executorch) project, so it needs the install above.
+        ${CONDA_RUN} cmake -S backends/mlx/examples/llm \
+          -B cmake-out/backends/mlx/examples/llm -DCMAKE_BUILD_TYPE=Release
+        ${CONDA_RUN} cmake --build cmake-out/backends/mlx/examples/llm \
+          -j$(( $(sysctl -n hw.ncpu) - 1 ))
+        RUNNER=cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf
+        if [ ! -x "${RUNNER}" ]; then
+          echo "Failed: runner not found at ${RUNNER}"
+          exit 1
+        fi
+        echo "::endgroup::"
+
+        echo "::group::Install LLM requirements"
+        ${CONDA_RUN} pip install -U "huggingface_hub[cli]<1.0"
+        ${CONDA_RUN} huggingface-cli login --token $SECRET_EXECUTORCH_HF_TOKEN
+        OPTIMUM_ET_VERSION=$(cat .ci/docker/ci_commit_pins/optimum-executorch.txt)
+        ${CONDA_RUN} pip install transformers "optimum-executorch @ git+https://github.com/huggingface/optimum-executorch.git@${OPTIMUM_ET_VERSION}"
+        if [ "${MODEL_ID}" = "google/gemma-4-E2B-it" ]; then
+          # Gemma 4 needs a newer Transformers than the CI-wide pin. Keep this
+          # on the same commit test-mlx-llm validated against.
+          GEMMA4_TRANSFORMERS_COMMIT=61461a7bcb458db7cf6eeea49678b9ab776a7821
+          ${CONDA_RUN} pip install -U "transformers @ git+https://github.com/huggingface/transformers.git@${GEMMA4_TRANSFORMERS_COMMIT}"
+        fi
+        echo "::endgroup::"
+
+        echo "::group::Export ${MODEL_NAME} off-graph"
+        ${CONDA_RUN} python -m executorch.backends.mlx.examples.llm.export_llm_hf \
+          --model-id "${MODEL_ID}" \
+          --output /tmp/${MODEL_NAME}_offgraph.pte \
+          --use-offgraph-cache \
+          --max-seq-len 1024 \
+          --dtype bf16 \
+          --qlinear 4w
+        echo "::endgroup::"
+
+        echo "::group::Run ${MODEL_NAME} off-graph inference"
+        # The cache geometry comes from the .pte; only capacity is given here.
+        TOKENIZER=$(${CONDA_RUN} python -c \
+          "from huggingface_hub import hf_hub_download; print(hf_hub_download('${MODEL_ID}', 'tokenizer.json'))")
+        OUTPUT=$(${RUNNER} \
+          --pte /tmp/${MODEL_NAME}_offgraph.pte \
+          --tokenizer "${TOKENIZER}" \
+          --chat "${CHAT}" \
+          --kv-max-capacity 1024 \
+          --prompt "What is the capital of France?" \
+          --max-new-tokens 50 2>&1)
+        echo "$OUTPUT"
+        if echo "$OUTPUT" | grep -iq "Paris"; then
+          echo "Success: 'Paris' found in output"
+        else
+          echo "Failed: Expected 'Paris' not found in output"
+          exit 1
+        fi
+        echo "::endgroup::"
@@ -0,0 +1,73 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# C++ runner for off-graph-cache HF models (export_llm_hf --use-offgraph-cache).
+# Unlike run_llm_hf.py (pybindings), this binds the cache via cache_key, so it
+# is the run path for off-graph .pte files.
+
+cmake_minimum_required(VERSION 3.24)
+project(mlx_run_llm_hf)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..)
+include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
+
+set(_common_include_directories ${EXECUTORCH_ROOT}/..)
+set(_json_include
+    ${EXECUTORCH_ROOT}/extension/llm/tokenizers/third-party/json/single_include
+)
+# MLXExecutor.h reaches schema_generated.h, which needs flatbuffers. The MLX
+# backend exposes that include under BUILD_INTERFACE, which does not reach a
+# separate find_package() project like this one.
+set(_flatbuffers_include ${EXECUTORCH_ROOT}/third-party/flatbuffers/include)
+
+list(APPEND CMAKE_FIND_ROOT_PATH ${CMAKE_CURRENT_BINARY_DIR}/../../../..)
+find_package(executorch CONFIG REQUIRED FIND_ROOT_PATH_BOTH)
+executorch_target_link_options_shared_lib(executorch)
+
+# As in examples/models/llama: point gflags_DIR at the copy the executorch build
+# configured, since a separate find_package() project cannot see it.
+set(gflags_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../../../third-party/gflags)
+find_package(gflags REQUIRED)
+
+set(link_libraries
+    executorch
+    extension_module
+    extension_tensor
+    extension_llm_cache
+    extension_llm_runner
+    extension_llm_sampler
+    gflags
+)
+
+if(NOT TARGET mlxdelegate)
+  message(FATAL_ERROR "mlx_run_llm_hf requires the MLX backend (mlxdelegate)")
+endif()
+list(APPEND link_libraries mlxdelegate mlx)
+executorch_target_link_options_shared_lib(mlxdelegate)
+
+# CPU kernels for the ops that stay outside the delegate (e.g. the cache
+# bookkeeping copy_ the HF export wrappers emit).
+if(TARGET optimized_native_cpu_ops_lib)
+  list(APPEND link_libraries optimized_native_cpu_ops_lib)
+  executorch_target_link_options_shared_lib(optimized_native_cpu_ops_lib)
+endif()
+
+list(APPEND link_libraries tokenizers::tokenizers)
+
+add_executable(mlx_run_llm_hf run_llm_hf.cpp)
+target_include_directories(
+  mlx_run_llm_hf PUBLIC ${_common_include_directories} ${_json_include}
+                        ${_flatbuffers_include}
+)
+target_link_libraries(mlx_run_llm_hf PUBLIC ${link_libraries})
+
+# The copy helper is gated on EXECUTORCH_BUILD_MLX, which the installed config
+# does not set; reaching here means mlxdelegate exists.
+set(EXECUTORCH_BUILD_MLX ON)
+executorch_target_copy_mlx_metallib(mlx_run_llm_hf)
diff --git a/backends/mlx/examples/llm/README.md b/backends/mlx/examples/llm/README.md
@@ -7,6 +7,7 @@ This example demonstrates how to export and run LLMs using the MLX delegate for
 - **Export**: Convert HuggingFace LLMs to ExecuTorch format with MLX delegate
 - **Quantization**: Optional INT4/INT8 weight quantization via TorchAO
 - **KV Cache**: Efficient KV cache implementation for autoregressive generation
+- **Off-graph KV Cache**: Optional runtime-owned cache, sized and configured per run instead of at export
 - **Custom Ops**: Uses `mlx::custom_sdpa` and `mlx::kv_cache_update` for optimal execution on MLX
 - **Pybindings**: Run inference using ExecuTorch Python bindings
 - **Gemma 4**: Text-only export and run flow supports processor-backed checkpoints such as `google/gemma-4-E2B-it`
@@ -23,6 +24,7 @@ pip install transformers optimum-executorch
 |--------|-------------|
 | `export_llm_hf` | Export LLMs using optimum-executorch pipeline, with optional custom MLX SDPA/KV cache |
 | `run_llm_hf` | Run exported models with token-by-token generation |
+| `run_llm_hf.cpp` | C++ runner; the run path for off-graph-cache exports, and runs in-graph ones too |
 
 For exporting via the ExecuTorch LLM pipeline (e.g. `examples/models/llama`), use `--mlx` to enable the MLX delegate.
 
@@ -54,6 +56,13 @@ python -m executorch.backends.mlx.examples.llm.export_llm_hf \
     --qlinear 4w \
     --qembedding 4w
 
+# Off-graph KV cache: the cache is owned by the runtime, not baked into the .pte
+python -m executorch.backends.mlx.examples.llm.export_llm_hf \
+    --model-id "unsloth/gemma-3-1b-it" \
+    --output gemma3_offgraph.pte \
+    --use-offgraph-cache \
+    --max-seq-len 1024
+
 # Gemma 4 text-only export
 python -m executorch.backends.mlx.examples.llm.export_llm_hf \
     --model-id "google/gemma-4-E2B-it" \
@@ -86,6 +95,12 @@ pip install -U "transformers @ git+https://github.com/huggingface/transformers.g
 | `--no-tie-word-embeddings` | `False` | Disable re-tying lm_head to embedding after quantization |
 | `--use-custom-sdpa` | `False` | Use MLX custom SDPA (`mlx::custom_sdpa`) |
 | `--use-custom-kv-cache` | `False` | Use MLX custom KV cache (`mlx::kv_cache_update`) |
+| `--use-offgraph-cache` | `False` | Use the off-graph KV cache (`kvcache::update_and_attend`); replaces the two flags above |
+| `--prefill-chunk-size` | `512` | Off-graph: tokens per prefill step, published as `get_prefill_chunk_size` for the runner. It is the largest single write, so a ring layer is sized `window + chunk - 1`; it may not exceed the sliding window or the exported sequence length |
+
+Off-graph exports keep no cache in the `.pte`, so the pybindings `run_llm_hf`
+cannot run them — use [`mlx_run_llm_hf`](#mlx_run_llm_hf-c) below, which builds
+the cache and binds it at load time.
 
 ---
 
@@ -123,6 +138,107 @@ python -m executorch.backends.mlx.examples.llm.run_llm_hf \
 
 ---
 
+## `mlx_run_llm_hf` (C++)
+
+Native runner for models exported with `--use-offgraph-cache`. That cache lives
+outside the graph and is handed to the backend by key at load time, which the
+pybindings `run_llm_hf` above cannot do. It also runs in-graph `.pte` files
+unchanged — omit `--kv-max-capacity` — so the same binary compares both cache
+paths. Greedy decode.
+
+### Build
+
+A standalone `find_package(executorch)` project, so ExecuTorch must be installed
+first:
+
+```bash
+cmake --preset mlx-release
+cmake --build cmake-out --target install -j$(( $(sysctl -n hw.ncpu) - 1 ))
+
+cmake -S backends/mlx/examples/llm -B cmake-out/backends/mlx/examples/llm \
+    -DCMAKE_BUILD_TYPE=Release
+cmake --build cmake-out/backends/mlx/examples/llm -j$(( $(sysctl -n hw.ncpu) - 1 ))
+```
+
+The binary lands at `cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf`.
+
+### Run
+
+```bash
+python -m executorch.backends.mlx.examples.llm.export_llm_hf \
+    --model-id unsloth/gemma-3-1b-it \
+    --output gemma3_offgraph.pte \
+    --use-offgraph-cache \
+    --max-seq-len 1024
+
+cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf \
+    --pte gemma3_offgraph.pte \
+    --tokenizer ~/.cache/huggingface/hub/models--unsloth--gemma-3-1b-it/snapshots/*/tokenizer.json \
+    --chat gemma \
+    --kv-max-capacity 1024 \
+    --prompt "What is the capital of France?" \
+    --max-new-tokens 50
+```
+
+`--chat` selects the instruct template — `llama3`, `gemma`, `gemma4`, or `0` for
+raw text. It matters: raw text confuses an instruct model into emitting turn
+markers, and using the wrong template invalidates a comparison between two
+`.pte` files.
+
+### Configuring the cache at run time
+
+Only the cache *geometry* is fixed at export — how many caches, their KV heads,
+head dims and windows, which the export publishes as constant methods the runner
+reads before building the cache. Everything else is chosen per run, with no
+re-export. The runner reports what it built:
+
+```
+[cache] off-graph seq | capacity=1024 initial=512 kv_dtype=BFloat16
+        26 layers: 4 flat + 22 ring(window 512)
+```
+
+`--kv-windows` overrides the attention pattern, repeating a comma-separated list
+over the caches (`0` = flat). The geometry each cache declared is untouched, so
+this cannot desync from the graph. For gemma-3-1b, whose 26 layers are 22
+sliding at 512 and 4 full:
+
+```bash
+# (omitted)                         26 layers: 4 flat + 22 ring(window 512)
+--kv-windows 512,512,512,512,512,0  # the same, spelling out gemma-3's 5:1 period
+--kv-windows 0                      # 26 layers: 26 flat
+--kv-windows 512                    # 26 layers: 0 flat + 26 ring(window 512)
+--kv-windows 512,256                # 26 layers: 0 flat + 13 ring(256) + 13 ring(512)
+```
+
+A ring layer allocates its whole `window + chunk - 1` slots up front, while a
+flat layer starts at `--kv-initial-capacity` and doubles as the sequence grows,
+so an all-flat cache can look smaller than a sliding one early in a run and
+larger later. Prefill runs in steps of the chunk size the export published, so
+the ring is bounded by the chunk rather than by the prompt: prefilling 24k
+tokens in one step would need 541 MiB of ring here, against 24 MiB at the 512
+default.
+
+### Options
+
+`--help` lists every flag with its default.
+
+| Option | Default | Description |
+|--------|---------|-------------|
+| `--pte` | *(required)* | Path to .pte file |
+| `--tokenizer` | *(required)* | Path to `tokenizer.json` |
+| `--prompt` | `The quick brown fox` | Input prompt |
+| `--max-new-tokens` | `50` | Tokens to generate, excluding the prompt |
+| `--temperature` | `0` | Sampling temperature; 0 is greedy argmax, which is what makes two `.pte` files comparable |
+| `--chat` | `llama3` | Chat template: `llama3`, `gemma`, `gemma4`, or `0` to disable |
+| `--kv-max-capacity` | `0` | Off-graph: history the cache may hold. Setting it selects the off-graph path |
+| `--kv-storage-dtype` | `bf16` | Off-graph: KV storage dtype (`bf16`, `fp16`, `fp32`) |
+| `--kv-initial-capacity` | `-1` | Off-graph: starting pool size; grows by doubling up to capacity |
+| `--kv-windows` | *(model's own)* | Off-graph: attention pattern override, e.g. `512` |
+| `--interactive` | `false` | Multi-turn chat on stdin; off-graph only |
+| `--warmup` | `false` | Run once before measuring, to absorb JIT and pool growth |
+
+---
+
 ## Architecture
 
 The `export_llm_hf` script uses optimum-executorch's `CausalLMExportableModule` by default. When custom flags are enabled, it uses `TorchExportableModuleWithStaticCache` from HuggingFace transformers, with optional MLX-specific replacements: