[MLX] C++ runner for HF LLMs with the off-graph KV cache

[kiymetakdemir]

Summary

This runner builds an MLXSequenceCache, installs it, and passes the cache key, so it's the run path for .pte files exported with --use-offgraph-cache. The cache's shape is read from the .pte metadata. The flags left are policy the model can't imply: --kv-max-capacity, --kv-storage-dtype, --kv-initial-capacity, --kv-max-write, --kv-windows.

Depends on #21680; the new CI job fails until that lands.

Files

• run_llm_hf.cpp — the runner: chat templates, greedy decode, benchmarking, and an interactive mode with /reset and /undo [N].
• CMakeLists.txt — standalone find_package(executorch) project.
• .github/workflows/mlx.yml — test-mlx-llm-offgraph for llama-1b, gemma3-1b and gemma4-e2b.

Test
CI builds the runner, then for llama-1b, gemma3-1b and gemma4-e2b exports off-graph and asserts the same "Paris" answer test-mlx-llm checks.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21683

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit fed3f79 with merge base 14326d9 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[metascroy]

Can you add something to the readme about this?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 6m 49s —— View job

---

Code review — MLX C++ runner for off-graph KV cache

Read run_llm_hf.cpp, CMakeLists.txt, the CI job, and the cache headers it binds against (cache.h, cache_registry.h, sequence_cache.h, backend_options.h). Overall this is clean, well-structured code — the rendezvous flow (build cache → install in registry → pass cache_key as a load option → load_method("forward")) matches the contract in cache_registry.h/backend_options.h exactly, the outer-scoped session/mlx_opts lifetimes are correct, and apply_window_override + valid(cfg) guard the config safely. A few issues worth addressing:

---

1. Missing mlx.metallib copy next to the runner (likely runtime break)

CMakeLists.txt never calls executorch_target_copy_mlx_metallib(mlx_run_llm_hf). The sibling qwen runner does exactly this (examples/models/qwen3_5_moe/CMakeLists.txt:90-93), and its CI even asserts the metallib landed next to the binary. MLX is statically linked here (add_library(mlx STATIC IMPORTED)), so at runtime MLX resolves mlx.metallib via dladdr() on the directory of the binary carrying MLX code — i.e. cmake-out/backends/mlx/examples/llm/. The install step only puts the metallib in lib/, not there. Unless every kernel this model touches is JIT-compiled (MLX_METAL_JIT=ON), the runner will fail to find mlx.metallib at first Metal dispatch.

if(TARGET mlxdelegate)
  executorch_target_copy_mlx_metallib(mlx_run_llm_hf)
endif()

Worth confirming against a real Metal run (the CI Paris check would catch it once #21680 lands). Fix this →

---

2. std::stoi on CLI args runs outside the try block (uncaught crash on bad input)

--kv-max-capacity, --max-new-tokens, --kv-initial-capacity, --kv-max-write, --warmup, --iters are all parsed with std::stoi(...) at run_llm_hf.cpp:240-263, before the try at line 270. A malformed value (e.g. --max-new-tokens abc) throws std::invalid_argument that nothing catches → std::terminate, no diagnostic. Either move these into the try, or wrap them the way parse_int_list already does. Given the file otherwise reports every error cleanly to stderr and returns 1, the abort is out of character. Fix this →

---

3. Interactive multi-turn drops the assistant's final token from history (subtle)

In the interactive loop (run_llm_hf.cpp:527-539), a generated token is only written to the cache when it becomes the input to the next step. The last token produced before is_stop/budget ends is printed but never fed back, so it's never committed to the KV cache. On the next user turn the model continues from history that's missing its own final token. This is the standard "last token isn't context" artifact and usually harmless, but for a /undo-style multi-turn tool it's a real (if minor) fidelity gap worth a one-line comment acknowledging it, or feeding the final token so history is exact.

---

Minor / nits

• run_llm_hf.cpp:259-261: the --interactive declaration is spliced into the middle of the comment block describing warmup/iters, so the "Multi-turn chat…" line reads as if it documents iters. Reorder so each comment sits above the flag it describes.
• run_llm_hf.cpp:443-447: step() throws on a failed execute, which in interactive mode tears down the whole session on a single transient failure. Acceptable for an example runner; flagging only in case interactive robustness matters.
• README: @metascroy asked for a README addition in a separate comment. Since your trigger was "review this code," I haven't implemented it — but the runner does deserve a short run_llm_hf (C++) section in backends/mlx/examples/llm/README.md alongside the existing pybindings row, covering the --kv-* flags and the export → build → run flow.

Nothing here blocks the design; #1 is the one I'd verify before merge.
· branch mlx-offgraph-hf-runner

[metascroy]

Why is kv-max-write a param?

Isn't it deriable from prefill-chunk-size?

[metascroy]

Can we use gflags?

[metascroy]

It's just a demo, but can we extract some of this cache config construction from pte metadata into a helper?