[MLX] C++ runner for HF LLMs with the off-graph KV cache by kiymetakdemir · Pull Request #21683 · pytorch/executorch

kiymetakdemir · 2026-08-07T23:58:22Z

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 · 2026-08-07T23:58:26Z

🔗 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.

❌ 3 New Failures

As of commit 589d63c with merge base 9958d39 ():

NEW FAILURES - The following jobs have failed:

Build Presets / apple (ios-simulator) / build (gh)
##[error]fatal: unable to access 'https://github.com/facebook/buck2-prelude.git/': SSL certificate problem: self signed certificate
Build Presets / apple (llm) / build (gh)
##[error]fatal: unable to access 'https://github.com/pytorch/cpuinfo.git/': SSL certificate problem: self signed certificate
Build Presets / apple (macos) / build (gh)
##[error]fatal: unable to access 'https://github.com/pybind/pybind11.git/': SSL certificate problem: self signed certificate

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

github-actions · 2026-08-07T23:59:13Z

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 · 2026-08-10T18:06:59Z

Can you add something to the readme about this?

metascroy · 2026-08-10T18:07:06Z

@claude review this code

claude · 2026-08-10T18:08:10Z

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 · 2026-08-10T21:54:15Z

+// Usage:
+//   run_llm_hf --pte <model.pte> --tokenizer <tokenizer.json> \
+//       [--kv-max-capacity N] [--kv-storage-dtype bf16|fp16|fp32] \
+//       [--kv-initial-capacity N] [--kv-max-write N] \


Why is kv-max-write a param?

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

There was no chunking, but added --prefill-chunk-size (default 512, overrideable) and dropped --kv-max-write.

Yeah, we can add a --prefill-chunk-size to export script with default value of 512.

Write the chunk size to metadata that can be read in the runner.

Error out during export if chunk size is bigger than ring size on sliding wndow.

Updated to this and if in runtime --kv-windows narrows window size it also narrows the chunk with it.

metascroy · 2026-08-10T21:58:02Z

+} // namespace
+
+int main(int argc, char** argv) {
+  const std::string pte = arg(argc, argv, "--pte", "");


Can we use gflags?

Changed to gflags.

metascroy · 2026-08-10T22:00:07Z

+    }
+
+    if (off_graph) {
+      cache::CacheConfig cfg{};


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

metascroy · 2026-08-11T16:52:57Z

+
+// argmax over the last position's vocab row of a [1, T, vocab] logits tensor,
+// reading whatever float dtype the op emitted.
+int64_t argmax_last(const ::executorch::aten::Tensor& logits) {


Don't we have a sampler/runner in ExecuTorch extension llm?

See llama runner and some of the other runners we have.

Changed it to reuse extension.

metascroy · 2026-08-11T16:54:08Z

+      const double pf_t = prompt_len / (pf / 1000.0);
+      const double dc_t = dc > 0 ? generated / (dc / 1000.0) : 0.0;
+      std::cout << "\n[iter " << iter << (measured ? "" : " warmup")
+                << "] prefill " << pf_t << " tok/s (" << prompt_len << " tok, "


I think we have something that collects stats on running?

Can you see some of the other C++ runners and reuse existign abstractions?

Switched to the shared Stats/print_report, removed --iters and changed --warmup to a bool to match the other runners and be able to use stats directly. Additionally adopted load_tokenizer and get_eos_ids abstractions.

…adata

kiymetakdemir · 2026-08-11T17:50:27Z

The MLX LLM jobs skip on fork PRs (they need HF secrets), is there anything I can do to run those jobs?

metascroy · 2026-08-11T18:10:18Z

@claude review this code

claude · 2026-08-11T18:11:15Z

Claude finished @metascroy's task in 5m 56s —— View job

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

Re-read run_llm_hf.cpp, CMakeLists.txt, export_llm_hf.py, README.md, and the cache interface (cache.h, cache_registry.h, sequence_cache.h) against the current branch. This is in much better shape than the first pass — the prior review items and Scott's comments are all addressed. A couple of small things remain, none blocking.

Resolved since last review

metallib copy (my Add support for quantized LeakyReLU #1) — CMakeLists.txt:73 now calls executorch_target_copy_mlx_metallib(mlx_run_llm_hf), gated by the mlxdelegate check above and with EXECUTORCH_BUILD_MLX set so the helper fires. Correct.
std::stoi crash (my Re-sync with internal repository #2) — gone with the move to gflags; the numeric flags are now DEFINE_int32/DEFINE_double, parsed safely by gflags. The only remaining hand-parsing (parse_int_list, /undo's std::stoll) is inside try/catch. Resolved.
final assistant token dropped (my Rename _pt2e to pt2e #3) — run_llm_hf.cpp:622 now commits the turn-end token so the turn is terminated in the cache. Good.
--kv-max-write derivable from chunk (Scott) — dropped; max_write now comes from get_prefill_chunk_size in the .pte metadata (read_kv_layout), export validates it against the sliding window (export_llm_hf.py:391,444), and --kv-windows narrows the chunk with the window (apply_window_override:243). Clean.
gflags (Scott) — done.
extract cache-config-from-metadata into a helper (Scott) — read_kv_layout / apply_window_override / const_int / const_ints. Done.
reuse ET extension abstractions (Scott) — now on load_tokenizer, get_eos_ids, Sampler + sample_from_logits, Stats + print_report, time_in_ms, safe_printf. Good reuse.
README (Scott) — the mlx_run_llm_hf (C++) section with build/run/flags is added.

Minor / nits (non-blocking)

mem_at_load isn't quite "weights-only" — run_llm_hf.cpp:437-438 calls it the "weights-only baseline, so the deltas below isolate the cache," but the cache is built at line 402 (before load_method), so get_active_memory() at line 438 already includes the initial pool. The +X MiB since load footprint then reports only cache growth from initial_capacity, not the full cache footprint. Either capture the baseline before CacheBuilderRegistry::build, or reword the comment/label so it doesn't imply the delta is the whole cache. Fix this →
Effectively-dead ternary in print_cache_summary — run_llm_hf.cpp:274-275 picks cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l], but read_kv_layout always pushes exactly n_layers entries and nothing shrinks the vector, so layers.size() always equals n_layers and the == 1 branch only differs when n_layers == 1 (where both branches are identical). It reads as if a single-entry config broadcasts to all layers, which isn't what happens here. If that broadcast semantics isn't intended, just index cfg.layers[l].
Budget-truncated turns still drop the last token (interactive) — the turn-end commit at line 622 fixes the normal stop case, but when a turn ends because budget is exhausted (not a stop token), the final next is neither printed nor fed back, so it's absent from history on the next turn. This is the documented "last token isn't context" artifact and is genuinely minor for a demo — noting only for completeness since Rename _pt2e to pt2e #3 from last time was about exactly this path.

Confirmed correct

Rendezvous flow (build cache → session.emplace → set kCacheKeyKey on mlx_opts → set_options → load_method("forward", …, &options_map)) matches the cache_registry.h/backend_options.h contract, and session/mlx_opts/options_map are outer-scoped so their lifetimes outlive the Module and the option view. Correct.
--kv-windows only rewrites per-layer policy and never touches the geometry the .pte declared, so it can't desync from the graph; apply_window_override narrows max_write down with the window and re-runs valid(cfg). Sound.
Interactive admission (can_extend(n+1), budget clamp, rewind/clear via control()) stays within capacity — no overflow path.

Nothing here blocks. #1 is the only item I'd bother touching before merge, and it's a comment/label fix.
· branch mlx-offgraph-hf-runner

metascroy · 2026-08-11T18:15:55Z

The MLX LLM jobs skip on fork PRs (they need HF secrets), is there anything I can do to run those jobs?

You should be able to add the "ciflow/mlx" label to trigger the MLX jobs. I just added it.

kiymetakdemir · 2026-08-11T19:59:18Z

I removed the kwargs check, HuggingFace forwards model-level arguments into attention (gemma 4 sends labels), so refusing anything unrecognised broke its export.

kiymetakdemir requested review from SS-JIA, digantdesai and metascroy August 7, 2026 23:58

kiymetakdemir self-assigned this Aug 7, 2026

kiymetakdemir requested review from kirklandsign and larryliu0820 as code owners August 7, 2026 23:58

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 7, 2026

metascroy reviewed Aug 10, 2026

View reviewed changes

metascroy reviewed Aug 11, 2026

View reviewed changes

kiymetakdemir added 2 commits August 11, 2026 10:01

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

c45d68e

Address review on the MLX off-graph runner

977de44

kiymetakdemir force-pushed the mlx-offgraph-hf-runner branch from dd06901 to 977de44 Compare August 11, 2026 17:41

Address review: reuse the shared LLM helpers, take the chunk from met…

1aa2e73

…adata

metascroy added the ciflow/mlx label Aug 11, 2026

metascroy approved these changes Aug 11, 2026

View reviewed changes

Stop rejecting unknown attention kwargs

589d63c

kiymetakdemir merged commit e7ca2b2 into pytorch:main Aug 11, 2026
251 of 262 checks passed

kiymetakdemir deleted the mlx-offgraph-hf-runner branch August 11, 2026 23:43

Uh oh!

Conversation

kiymetakdemir commented Aug 7, 2026

Uh oh!

pytorch-bot Bot commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 3 New Failures

Uh oh!

github-actions Bot commented Aug 7, 2026

This PR needs a release notes: label

Uh oh!

metascroy commented Aug 10, 2026

Uh oh!

metascroy commented Aug 10, 2026

Uh oh!

claude Bot commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

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

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

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

Minor / nits

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kiymetakdemir Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

metascroy Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kiymetakdemir commented Aug 11, 2026

Uh oh!

metascroy commented Aug 11, 2026

Uh oh!

claude Bot commented Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

Resolved since last review

Minor / nits (non-blocking)

Confirmed correct

Uh oh!

metascroy commented Aug 11, 2026

Uh oh!

kiymetakdemir commented Aug 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pytorch-bot Bot commented Aug 7, 2026 •

edited

Loading

This PR needs a `release notes:` label

claude Bot commented Aug 10, 2026 •

edited

Loading

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

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

kiymetakdemir Aug 11, 2026 •

edited

Loading

metascroy Aug 11, 2026 •

edited

Loading

claude Bot commented Aug 11, 2026 •

edited

Loading