Skip to content

[Native] Run off-graph KV cache packages in the LLM runner - #24638

Draft
metascroy wants to merge 1 commit into
offgraph-kv-ggml-enginefrom
offgraph-kv-runner
Draft

metascroy wants to merge 1 commit into
offgraph-kv-ggml-enginefrom
offgraph-kv-runner

Conversation

@metascroy

@metascroy metascroy commented Nov 5, 2026 •

Copy link
Copy Markdown
Contributor

Stack

Off-graph KV cache for the native ggml engine, through serving Muse Glimmer on it. Review bottom-up; each PR builds and passes its tests on its own.

Summary

run_hf_llm_ggml (via LlmRunner) can now run packages exported with --offgraph_cache.

  • Cache shape: the runner reads it from the graph. Each update_and_attend call's layer_id and K shape give one flat layer, so the .ptn format doesn't change.
  • Setup: it builds the engine's single-sequence cache with CacheFactory::build(backend_id, kind::kSingle, ...) and publishes it with an InstallGuard. It passes the registry key to create_context_with_options, and keeps the cache alive longer than the context.
  • New flags:
    • --kv_capacity: tokens; defaults to --max_seq_length.
    • --kv_dtype: fp32 or fp16.
    • --kv_initial_capacity: rows allocated up front.
  • --verify_samples clears the cache between samples instead of recreating the context.
  • run_llm takes the engine's cache backend id. run_hf_llm_ggml registers the ggml cache builders and passes NativeGgml.
  • --kv_capacity, --kv_dtype and --kv_initial_capacity are errors on an in-graph package instead of being silently ignored.
  • extension_llm_cache is now also built when EXECUTORCH_BUILD_NATIVE_LLM_RUNNER is on.
run_hf_llm_ggml --device gpu --package llama1b_q4km_offgraph.ptn \
  --tokenizer tokenizer.model --kv_capacity 2048 --kv_dtype fp16 --prompt ...

Test plan

On an Apple M4 Max:

  • Tiny 2-layer random-weight off-graph export, CPU, --verify_samples golden --kv_initial_capacity 1: 0 calls failed, max error about 1.5e-5.
  • Llama 3.2 1B Q4_K_M off-graph, Metal, --verify_samples golden: 0 calls failed (atol 2, rtol 0.05; max error 0.05-0.14).
  • A chat-templated prompt on the off-graph and in-graph packages: same output ("The capital of France is Paris.").

@pytorch-bot

pytorch-bot Bot commented Nov 5, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 147 New Failures, 1 Cancelled Job, 7 Unrelated Failures, 27 Unclassified Failures

As of commit 24b19de with merge base c7f7931 (image):

NEW FAILURES - The following jobs have failed:

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

CANCELLED JOB - The following job was cancelled. Please retry:

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

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

@meta-cla 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 Nov 5, 2026
@metascroy
metascroy force-pushed the offgraph-kv-ggml-engine branch from 2f88227 to ab17a5f Compare November 5, 2026 22:29
@metascroy
metascroy force-pushed the offgraph-kv-ggml-engine branch from ab17a5f to 8128a44 Compare November 11, 2026 18:05
Comment on lines +71 to +79
DEFINE_int32(
kv_capacity,
0,
"Off-graph KV cache capacity in tokens; 0 uses --max_seq_length. Only "
"for packages exported with --offgraph_cache.");
DEFINE_string(
kv_dtype,
"fp32",
"Off-graph KV cache storage: fp32 or fp16.");

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: MLX's runners call this --kv_max_capacity

When a method calls kvcache::update_and_attend, the runner derives the cache
geometry from the graph, builds the engine's single-sequence cache through
the CacheFactory, publishes it with an InstallGuard, and passes its registry
key to the engine context. New flags --kv_capacity, --kv_dtype, and
--kv_initial_capacity size the cache. Sample verification clears the cache
between samples instead of recreating the context. run_hf_llm_ggml registers
the ggml cache builders.

This branch had an error being deployed

1 failed deployment
cadence — 24b19de0 Deployed Nov 14, 2026 by metascroy via hifi-build / hifi4 #38329
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants