Skip to content

[Native] Serve kvcache::update_and_attend in the ggml engine - #24637

Draft
metascroy wants to merge 1 commit into
offgraph-kv-engine-optionsfrom
offgraph-kv-ggml-engine
Draft

[Native] Serve kvcache::update_and_attend in the ggml engine#24637
metascroy wants to merge 1 commit into
offgraph-kv-engine-optionsfrom
offgraph-kv-ggml-engine

Conversation

@metascroy

@metascroy metascroy commented Nov 5, 2026 •
edited
Loading

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

The native ggml engine can now run graphs exported with an off-graph KV cache (kvcache::update_and_attend).

Cache: GgmlSequenceCache (new, GgmlCache.{h,cpp})

  • It's a SequenceCache, so it uses the neutral sequence controller (SequenceControl/SequencePlanner). It adds K/V storage per layer, laid out [head_dim, rows, n_kv_heads] in the engine backend's default buffer type.
  • Storage starts at the config's initial_capacity rows and doubles toward capacity, copying existing rows across. Short sequences don't pay for the full context.
  • Each allocation is a reference-counted Storage. An executable holds the storage its graphs point into until it rebinds, so a regrow by another executable never leaves dangling tensors behind.
  • Like SequenceCache, it isn't thread-safe. Contexts bound to one cache must be serialized; create_context_with_options documents this.
  • register_ggml_cache_builders() registers it with the CacheFactory under backend id NativeGgml, kind single. Registration is an explicit call rather than a static initializer, so a static link can't drop it.

Engine

  • GgmlEngineHost implements create_context_with_options. The llm_cache_registry_key option resolves the cache from the CacheRegistry. A missing key or an unknown option throws.

  • update_and_attend lowers to ggml_set_rows into the layer's cache, a view of the filled prefix, and ggml_flash_attn_ext with an F16 causal mask. The rows and mask are per-step uploads shared by all layers.

  • How many cache rows to attend is part of the graph cache key. It's max(position) + 1, rounded up to 256 and capped at the allocated rows. HostIntEvaluator learns select.int so it can evaluate HF's position_ids[0].

  • Before each call, the executor reads the positions once on the host, then:

    1. checks they're consecutive, fit the capacity, and start no later than the sequence's length, at a position the sequence can rewind to. A step starting earlier rewinds to its start once it has run, as re-running the same positions (--warmup_prefill) does; one starting past the length is rejected, since it would attend rows nothing wrote;
    2. grows the cache, and on a regrow rebinds the storage and drops cached graphs built over the old one;
    3. plans every layer, then fills the rows and mask from the plan.

    It rewinds and commits only after the call succeeds, so a call that fails before its graph runs (a host-side index check, growing the cache, building the graph) leaves the sequence unchanged. One that fails while running may have written rows, so the sequence drops its positions from the step's start.

    An F16 step written into an F32 cache is converted first: ggml_set_rows doesn't convert F16 to F32 on every backend (CUDA aborts, Metal falls back to CPU).

    The cache builder reports a geometry or config the engine can't serve through CacheFactory::build's error result rather than an exception.

Limits: flat layers only (sliding-window layers are rejected), batch 1, fp32 or fp16 storage.

Test plan

New test_ggml_engine.cpp tests. Each runs update_and_attend against a naive softmax attention over everything the cache should hold:

  • OffGraphAttentionMatchesReferenceAsTheCacheGrows:
    • a 3-token prefill, then decode steps, in two contexts sharing one cache;
    • storage starts at one row, so it regrows during the prefill and the decode;
    • with 0, 1, and 4 cached graphs, and fp32 and fp16 storage.
  • OffGraphAttentionReplaysFromEmptyAfterClear, including an executable whose binding went stale because another one regrew the cache.
  • OffGraphAttentionContinuesFromARewind.
  • OffGraphStepsStartAtOrBeforeTheLength: re-running a step rewinds; skipping past the length is rejected.
  • FailedOffGraphStepLeavesTheSequence: a step starting before the length that fails on a host-side row check leaves the sequence where it was. It fails if the rewind happens before the graph is built.
  • RejectsOffGraphStepsPastCapacityOrNotConsecutive.
  • OffGraphAttentionNeedsAMatchingCache: no key, an unknown key, an unknown option, a layer out of range, and K/V shaped unlike the cache.

Checked that the tests catch real bugs: skipping the row copy on regrow fails 3 of them, dropping the causal mask fails 4, and dropping the start-position check fails 2.

native_ggml_engine_test: 43 passed. GpuUnavailableOnCpuVariant fails as it does on main, because this build has Metal enabled.

@pytorch-bot

pytorch-bot Bot commented Nov 5, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

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

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

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

As of commit 80c68e9 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.

This was referenced Nov 5, 2026
Draft
Draft
Draft
Draft
Draft
@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-engine-options branch from 4b5de65 to 0c97fdf Compare November 5, 2026 22:29
@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
@metascroy
metascroy force-pushed the offgraph-kv-engine-options branch from 0c97fdf to 4ec9257 Compare November 11, 2026 18:05
@kiymetakdemir

kiymetakdemir commented Nov 11, 2026

Copy link
Copy Markdown
Contributor

@claude /pr-review

This was referenced Nov 12, 2026
Draft
Draft
Draft
Draft
Draft
Draft
Draft
Closed
Closed
Draft
Draft
Draft
Draft
kiymetakdemir
kiymetakdemir approved these changes Nov 12, 2026
View reviewed changes

@kiymetakdemir kiymetakdemir left a comment •
edited
Loading

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.

Most of the Claude review points look like nits to me, but they'd make the design safer and are easy to address such as

  • check that positions have one component per token before the consecutive check,
  • check can_rewind(start) in checked_offgraph_step and rewind just before run(), so a step that fails before running leaves the length unchanged.

My question is whether start < length should be allowed, and if so, which behavior the backends should share: rewinding to start as this does, or the MLX cache's in-place overwrite, which keeps the length.

Comment thread backends/native/runtime/engine/ggml/GgmlExecutor.cpp
", past the sequence's length of " +
std::to_string(kv_cache_->length()));
}
if (start < kv_cache_->length() &&

@kiymetakdemir kiymetakdemir Nov 12, 2026

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.

When should a step be allowed to start before the length? Here the only caller that does it is --warmup_prefill, which re-runs the prompt at 0. If that's the only case, the runner could call kv->control->clear() after warmup, and the engine could require start == length(). DFlash calls rewind() explicitly to drop rejected draft tokens.

If we do want to allow it, should it change the length? MLX overwrites in place and keeps the length. Here the step rewinds to its start, so the length becomes start + tokens. Rewinding to start as this does seems more correct than MLX's in-place overwrite.

This was referenced Nov 13, 2026
Draft
Draft
Draft
Draft
Draft
@metascroy
Serve kvcache::update_and_attend in the native ggml engine
80c68e9 
The ggml engine now runs graphs exported with an off-graph KV cache.

- GgmlSequenceCache is the ggml SequenceCache: the neutral sequence
  controller plus K/V storage per layer in the engine's backend buffer.
  Storage starts at the config's initial_capacity rows and doubles toward
  capacity, so short sequences do not pay for the full context.
- register_ggml_cache_builders() registers it with the CacheFactory under
  backend id NativeGgml, kind single.
- A context created with the llm_cache_registry_key option resolves the cache
  from the CacheRegistry.
- update_and_attend lowers to ggml_set_rows into the cache and
  flash_attn_ext over the filled prefix, with per-step rows and causal mask
  uploads. Before each call the executor checks positions, grows the cache,
  rebinds storage (dropping cached graphs on a regrow), and plans every
  layer; it commits after the call succeeds.

Flat layers, batch 1, and fp32/fp16 storage only.
@metascroy
metascroy force-pushed the offgraph-kv-engine-options branch from 4ec9257 to b2aebff Compare November 14, 2026 00:18
@metascroy
metascroy force-pushed the offgraph-kv-ggml-engine branch from 8128a44 to 80c68e9 Compare November 14, 2026 00:18

This branch had an error being deployed

1 failed deployment
cadence — 80c68e90 Deployed Nov 14, 2026 by metascroy via hifi5s-build / hifi5s #38334
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@kiymetakdemir kiymetakdemir kiymetakdemir approved these changes
@larryliu0820 larryliu0820 Awaiting requested review from larryliu0820 larryliu0820 will be requested when the pull request is marked ready for review larryliu0820 is a code owner

Assignees

No one assigned

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

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@metascroy @kiymetakdemir