Expose committed session cloning and build prefix reuse above it

[mergennachin]

Expose Session::clone_async(upto) for independently writable committed prefixes. A caller-owned PrefixCache handles token matching and LRU eviction; capture_prompt(...).wrap(...) requests the prompt clone before the first user callback, and caller-thread collect() inserts it after generation. Capture binds to session identity and survives moves without retaining source ownership. Lookup tries each snapshot only at its longest matching boundary. The MLX example uses this helper and forwards the final prompt token for fresh logits. The cache factory API stays unchanged; session limits are checked through the constructed cache's BatchControl::max_seqs().

Review the Session/Runner lifecycle first, then ModuleExecutor/backend clone contracts, prefix policy, and MLX integration.

214 shared native tests pass on Apple M1 Pro, including 14 capture tests for timing, reuse, failures, handle lifetime, and moved/reassigned/destroyed sources. All five native targets and the standalone MLX batched runner build; lint passes. GPU execution was not repeated for this lifetime change; the following GPU/model results are from ddf92b40d1.

All 49 MLX tests pass. Gemma 3 1B validation covers BF16/FP16, full/selected logits, both layouts, physical ring wrap, branches, closure, and concurrency: 16 smoke groups, 88 cache/work checks, and 44 cold export comparisons pass. Warm repeats forward 1 of 23/1280 prompt tokens. Normal cold/warm initial logits exceed tolerance in 72/88 cases. The public header recommends opt-in greedy experiments until non-greedy sampling is validated; even greedy token parity is not guaranteed. Clone controls with matching schedule/geometry are byte-exact. A 23-token control also gives byte-exact cold/split/replay/clone logits in both dtypes with one-token prefill. Isolated probes identify shape-dependent MLX projection and attention arithmetic; precision promotions did not resolve the normal cold/warm gap. The four existing exports were reused and hash-verified.

Authored with OpenAI Codex.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 183 Pending

As of commit 59d22ca with merge base 026ca3f ():
💚 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.

[Copilot]

🔵 Needs a closer look

Unresolved MLX logit drift, allocation risk, and cache-hit metric distortion require changes and human review.

Pull request overview

Adds an opt-in, bounded prefix cache for off-graph batched LLM sessions, reusing immutable snapshots and replaying the final requested token for logits.

Changes:

• Adds longest-prefix lookup, LRU eviction, cloning, and retention handling.
• Integrates cache admission, publication, replay, capacity, lifecycle, and metrics.
• Adds MLX support, examples, documentation, tests, and CI coverage.

File summaries

File	Reviewed changes and final notes
extension/llm/cache/test/prefix_cache_test.cpp	Tests prefix lookup, LRU policy, retention, and snapshot ownership.
extension/llm/cache/test/prefix_cache_attention_test.cpp	Tests attention correctness across cache layouts.
extension/llm/cache/test/CMakeLists.txt	Registers shared cache tests.
extension/llm/cache/prefix_cache.h	Implements bounded prefix snapshots and lookup.
extension/llm/cache/cache.h	Defines the backend-independent clone contract.
extension/llm/batching/test/runner_test.cpp	Tests reuse, replay, accounting, cancellation, and shutdown.
extension/llm/batching/runner.h	Documents opening-prompt reuse semantics.
extension/llm/batching/runner.cpp	Adds cache admission, publication, replay, and metrics. Moderate (1 vote): prefill throughput should exclude cached tokens or include restoration time.
extension/llm/batching/README.md	Documents prefix-cache configuration and behavior.
extension/llm/batching/module_executor.h	Exposes prefix-cache configuration and hooks.
extension/llm/batching/module_executor.cpp	Integrates cache construction and capacity sizing. Critical (1 vote): validate sequence limits before allocation to avoid OOM instead of returning InvalidArgument.
extension/llm/batching/metrics.h	Adds cached-token metrics.
extension/llm/batching/metrics.cpp	Reports cached-token metrics.
extension/llm/batching/executor.h	Adds optional executor prefix hooks.
backends/mlx/test/mlx_prefix_cache_test.cpp	Tests MLX cache reuse and attention. Moderate (1 vote): add a production-model regression or gate caching until reported logit drift is explained.
backends/mlx/test/CMakeLists.txt	Registers MLX prefix-cache tests.
backends/mlx/examples/llm/run_llm_batched.cpp	Adds cache flags and repeated prompt rounds.
backends/mlx/examples/llm/README.md	Documents MLX usage and smoke tests.
.github/workflows/mlx.yml	Adds expanded MLX cache test coverage.

Review details

Suppressed comments (2)

backends/mlx/test/mlx_prefix_cache_test.cpp:114

• This oracle only exercises a hand-written two-layer attention model. The PR's own MLX validation reports cold-versus-warm logit mismatches on the real model even in replay controls without cloning, so this assertion can pass while the production cache changes logits. Please add a production-model regression that isolates the drift, or keep the cache path gated until that numerical discrepancy is explained.

    EXPECT_TRUE(allclose(got, expected, 1e-2f));

extension/llm/batching/runner.cpp:1285

• Recording the cache hit here makes GenerationMetrics::prefill_tokens_per_sec() use the full n_prompt_tokens, including these cached tokens, while prefill_span_us() starts at the first executed batch and therefore excludes prefix restoration. A large cache hit will consequently report cached work as prefill throughput (and can make the rate arbitrarily high). Compute this rate from the uncached prompt tokens, or add a separate metric whose denominator includes restoration time.

      request.generation.m.n_cached_prompt_tokens = *reused;

• Files reviewed: 19/19 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[metascroy]

Could we expose cloning as a general session operation and build prefix caching on top of that, rather than adding reuse_prefix() and cache_prefix() to Executor?

Concretely, Session::clone_async(upto) would route through the runner to Executor::clone(source, upto), with ModuleExecutor delegating to BatchControl::seq_clone(). The result would be a new independently writable session containing the committed prefix, with backend retention checks and resource limits determining whether cloning succeeds.

The prefix cache would then own retained snapshot sessions and the token-matching/LRU policy. A lookup would find the longest usable match and return a clone plus the matched-token count, leaving at least the final prompt token to be forwarded. On a miss, the caller opens a fresh session; on eviction, the cache releases its retained session.

Thoughts?

[Copilot]

🟡 Changes recommended

Unresolved correctness, ordering, and test deadlock issues block approval.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

extension/llm/batching/runner.cpp:780

• A clone enqueued from an output callback is not processed before the next forward: process_pending_commands_() runs only at the top of run_(), while a successful execute_one_batch_() immediately continues around that call. The continuation can therefore advance/wrap the source before this CloneCommand runs, causing prompt snapshots to capture too late or be refused by sliding-window retention, contrary to the documented callback ordering. Drain clone/close commands between the completed forward and scheduling its continuation, or otherwise establish that ordering.

            [this](CloneCommand& clone) { process_command_(std::move(clone)); },

• Files reviewed: 18/18 changed files
• Comments generated: 2
• Review effort level: Lite

[mergennachin]

The numerical investigation at ddf92b40d1 now isolates MLX arithmetic differences in the 23-token BF16/FP16 selected-logits, batched-sequence reproducer. All traced/untraced and archived output checks pass byte-for-byte. The first difference is Q projection (op 35) on identical input and weights. Isolated attention replay reproduces the larger first-layer discrepancy: rounding QK scores and softmax probabilities to the model dtype reconstructs every unfused output value in both dtypes. Matched clone/original replay remains byte-exact.

A native control processing the entire prompt one token per forward gives byte-identical cold, split, original replay, and clone replay logits: all 12 pairwise comparisons pass across BF16/FP16, over 262,144 logits per row. All 32 main/preparation schedules and cache/work counters are checked, including reuse of 22 tokens and fresh execution at position 22. Ordinary multi-token prefill still reproduces the original 0.75 BF16 / 0.078125 FP16 maximum error. This supports schedule-dependent numerical drift; no prefix-policy or clone-state defect is demonstrated by this repro.

Uniform low-precision attention, FP32 attention, and FP32 attention plus projections were also tested; none resolves the ordinary cold/warm tolerance failure. This does not invalidate the local attention diagnosis, but full-model error growth is not completely traced. Serial prefill is a diagnostic control, not a production workaround.

The broader validation remains 251 native tests, 16 smoke groups, 88 cache/work checks, and 44 cold full/selected comparisons passing, with 72/88 ordinary cold/warm initial-logit comparisons outside the unchanged tolerances. The matching-schedule result covers one short prompt and one initial prediction; it is not a general numerical-equivalence claim. Non-greedy sampling remains untested and prefix reuse remains disabled by default (--prefix_cache_entries=0). The PR summary now records these findings.

[Copilot]

🔵 Needs a closer look

A moderate close-versus-executor-failure lifecycle issue remains, and the broad backend/MLX changes warrant human review.

Review details

Suppressed comments (1)

extension/llm/batching/runner.cpp:760

• The new !status->open check is only reached by reap_cancelled_ between calls to execute_one_batch_. If a session is destroyed while its in-flight executor call returns false (or produces an incomplete output), execute_one_batch_ completes the generation as Failed and poisons the session before this check runs, even though Session destruction promises cancellation. Reclassify closed sessions before the failure/malformed-output completion paths as Cancelled (and avoid poisoning them), and add a close-vs-executor-failure regression test.

    if (generation &&
        (generation->state->cancelled.load() ||
         !entry.second.status->open.load(std::memory_order_acquire))) {

• Files reviewed: 19/19 changed files
• Comments generated: 0 new
• Review effort level: Lite

[metascroy]

Since correct snapshot capture depends on requesting a clone from the first output callback, could PrefixCache provide a small helper for that pattern? Ideally it would wrap the generation callback, request the clone once without blocking, and expose a caller-thread collection/insertion step. That would keep users from having to reproduce the timing, future ownership, and failure-handling logic in each integration.

[mergennachin]

Implemented in 59d22ca57d:

auto capture = prefix_cache.capture_prompt(session, prompt_tokens);
auto generation = session.generate_async(
    suffix, config, capture.wrap(on_update));
generation.wait();
capture.collect();

The wrapper requests one clone before the first nonempty user callback, without waiting. Collection inserts best-effort on the caller thread. Capture binds to the original session ID and status, so moving or reassigning the source wrapper is safe; it does not own the source or delay closure. The MLX example and docs use this flow.

All 214 shared tests pass, including 14 capture cases and three source-lifetime regressions. All five native targets and the standalone MLX runner build; lint passes. GPU execution was not repeated. The header also states lookup's one-boundary-per-snapshot behavior and the MLX cold/warm parity caveat, including the lack of a greedy-token guarantee.

[mergennachin]

@claude Review this PR

[claude]

Claude finished @mergennachin's task in 6m 8s —— View job

---

PR Review: Expose committed session cloning and build prefix reuse

I reviewed the current tree (post-dfe955e) in review order: Session/Runner lifecycle → clone contract → prefix policy → MLX. The refactor asked for in the API-review thread landed cleanly, and it retires most of the earlier Copilot findings. Below is what I found on the code as it stands now.

Architecture — the refactor holds up

Moving prefix reuse out of Runner/Executor and onto a caller-owned PrefixCache over Session::clone_async(upto) is the right call, and it resolves several stale review comments that were written against the old design:

• Cache-hit metric distortion (old runner.cpp:1285 / n_cached_prompt_tokens): gone. metrics.h no longer carries cached-token fields, so there's no prefill_tokens_per_sec denominator skew anymore.
• Pre-allocation OOM (old module_executor.cpp): the separate sequence-limit metadata was removed; create() now builds the cache and does the post-construction BatchControl::max_seqs() check (module_executor.cpp:384-392). Consistent with the layout owning its own limit.
• CloneCommand ordering (old runner.cpp:780): the run_() loop is process_pending_commands_() → execute_one_batch_() → continue, and continue re-enters at the loop top, so a clone enqueued from an output callback is drained before the next execute_one_batch_(). Critically, no forward runs between the first-callback enqueue and the clone, so the source hasn't advanced — this is what keeps a sliding-window clone of [0, prompt_len) from being refused by ring retention. The capture_prompt().wrap() "clone from the first callback" pattern is load-bearing for SWA, and it's correct.

clone failure/shutdown paths in process_command_(CloneCommand) (runner.cpp:838-889) look leak-free: on a non-published clone the code erases the registered record and calls close_session, and the issued_session_ids_ set only ever grows (its job is uniqueness detection, not reclamation).

Remaining issues

1. PrefixCache::lookup doesn't retry shorter boundaries within a snapshot (minor correctness-vs-doc gap) — prefix_cache.h:215-224
The doc promises "longest usable token match," and refusal handling says "retention/resource refusal may make a shorter match the best usable one." But the loop tries each candidate only at its single longest boundary. If snapshot A matches 20 tokens and its clone is refused under resource pressure, A is abandoned even though a clone at, say, 8 tokens might succeed — and lookup then returns a shorter match from a different snapshot (or a miss) while A held a usable longer prefix. This is the one Copilot comment on the current file that still stands. It's low-severity (any returned match is still valid and you always forward the final token), but the behavior doesn't match the contract's wording. Either implement per-snapshot boundary backoff and compare attempts globally, or soften the doc to "longest match whose full boundary clones."

2. PromptCapture::State holds const Session& source by reference — prefix_cache.h:102
Session is movable, so if a caller moves the borrowed session (e.g. reseats it in a container, or reassigns) between capture_prompt(...) and generation completion, the stored reference dangles and wrap()'s state->source.clone_async(...) runs on the engine thread against freed storage. The header does document "keep source alive at the same address," so this is a documented footgun rather than a bug, but given the handle is explicitly move-friendly and outlives the call, a sentence pinning down "must not move the source Session either" (or storing what's actually needed — the RunnerImpl + SessionId — rather than a Session&) would harden it. Fix this →

3. MLX warm/cold logit drift is the material risk — currently mitigated only by "disabled by default"
Your investigation at ddf92b40d1 is thorough and the conclusion is credible: clone/replay with a matching schedule is byte-exact, so the ~0.75 BF16 / 0.078 FP16 gap (72/88 cases) is schedule-dependent MLX projection/attention arithmetic, not a prefix-policy or clone-state defect. That's a strong diagnosis. The residual concern is behavioral, not the diagnosis: a warm prefix reuse changes the schedule the initial logits are produced under, so under non-greedy sampling a cache hit can sample a different token than the cold path — i.e. reuse isn't purely a work-saving optimization. It's gated behind --prefix_cache_entries=0 by default and non-greedy is untested, which is the right posture for an experimental helper, but I'd make the equivalence caveat impossible to miss: keep the "cold/warm parity is not guaranteed under this backend; greedy-only until traced" note directly in the PrefixCache header (not only in the README/PR body), since that's what an integrator reads first.

Smaller notes

• README.md (batching) and backends/mlx/examples/llm/README.md document the capture_prompt(...).wrap(...) / collect() flow well; the example matches. Good.
• PrefixCache non-thread-safety and "never call lookup() from a Runner callback" (it blocks on .get()) are clearly documented at the class head. Worth keeping that prominent since collect() (caller thread) and lookup() (caller thread) both mutate entries_ and must be externally serialized.
• The ET_HAS_EXCEPTIONS guarding is consistent throughout and bad_alloc is contained on every allocation path in the cache — nice attention to the no-exceptions build.

Verdict

No blocking bug in the exposed cloning path or the prefix policy on the current tree; the lifecycle and clone-contract reasoning is sound and matches the tests you describe. Item 1 is a doc/impl mismatch worth tightening, item 2 is a lifetime footgun worth pinning down in the contract, and item 3 is the real risk to flag for anyone who turns the cache on — the disabled-by-default + greedy-only posture is appropriate for landing this as experimental. Since @metascroy has already approved, these are refinements rather than merge blockers.

--- • branch llm-prefix-cache-swa

[Copilot]

🔵 Needs a closer look

MLX validation reports substantial cold/warm initial-logit mismatches requiring final human review.

Review details

• Files reviewed: 19/19 changed files
• Comments generated: 0 new
• Review effort level: Lite

[mergennachin]

@claude Re-review now

[claude]

Claude finished @mergennachin's task in 1m 47s —— View job

---

Re-review: prefix cache capture refactor

I re-read the current tree (commit 59d22ca) in review order — Session/Runner clone lifecycle → ModuleExecutor clone contract → PrefixCache policy/capture → MLX. My three prior items map cleanly onto this revision:

Prior items — status

Item 2 (dangling const Session&) — resolved. The footgun is gone. PromptCapture::State no longer stores a Session&; it holds a request_clone functor produced by Session::make_clone_request(upto) (runner.cpp:579-596), which captures impl, the SessionId, status, and upto by value. Moving or reassigning the source Session between capture_prompt(...) and collect() can no longer dangle, and the functor re-checks status->open before routing to clone_async. The header now states this directly: "Capture binds to the session identity, so the source Session may move… destroying its owner still requests closure" (prefix_cache.h:124-128). This is the stronger of the two fixes I suggested (store what's needed rather than a reference) — good.

Item 3 (MLX cold/warm parity caveat visible to integrators) — resolved. The equivalence caveat is now in the PrefixCache class header itself, not just the README/PR body: "On MLX, cold/warm logits may differ as prefill shapes change. Keep MLX reuse opt-in and greedy-only until non-greedy sampling is validated; even greedy token parity is not guaranteed" (prefix_cache.h:42-45). That's what an integrator reads first, which was the ask.

Item 1 (lookup "longest usable match" doc/impl gap) — resolved by doc, as I offered. Rather than implementing per-snapshot boundary backoff, the contract was softened to match the code: "Each snapshot is tried once at its longest match; refusal skips it without retrying shorter boundaries" (prefix_cache.h:189-192). lookup() now collects all matching candidates, sorts by (match length, recency), and tries each at its single longest boundary (prefix_cache.h:194-242). Doc and implementation now agree. This is the acceptable resolution — any returned match is still valid and the final token is always forwarded. (The Copilot comment on this is now stale.)

Spot checks on the new capture path

• Clone-before-first-forward ordering still holds. run_() is process_pending_commands_() → reap_cancelled_() → execute_one_batch_() → continue (runner.cpp:723-734); wrap() enqueues the clone on the first non-empty update (prefix_cache.h:66-77) and the source hasn't advanced between that enqueue and the drain, so a [0, prompt_len) snapshot isn't refused by ring retention. Load-bearing for SWA, and intact.
• collect() / no-exceptions build. collect() moves state_ out first, so a second call is a safe no-op returning false (prefix_cache.h:87-91); bad_alloc is contained on every allocation path in capture, insert, and lookup. insert enforces snapshot.position() == tokens.size() (prefix_cache.h:165), so a partially-committed clone can't be admitted.
• CloneCommand leak paths (runner.cpp:857-909) remain clean: on a non-published clone the code erases the registered record and calls close_session; issued_session_ids_ only grows (uniqueness detection, not reclamation).
• ModuleExecutor::clone (module_executor.cpp:472-491) guards upto, max_session_tokens_, max_sessions_, and next_session_ == 0 before seq_clone, and publish_session's SequenceGuard releases the seq on any publish failure. The max_seqs() post-construction check is back per the API-review thread; the pre-allocation OOM concern is documented as a known limitation (large CPU cell metadata may allocate before rejection; GPU KV is lazy).

Remaining note (unchanged, non-blocking)

The MLX warm/cold logit drift (72/88 ordinary cases outside tolerance at ddf92b40d1) is still the material behavioral risk, and the diagnosis — schedule-dependent MLX projection/attention arithmetic, not a clone-state or prefix-policy defect — is well-supported by the byte-exact matched-schedule and serial-prefill controls. Reuse remains disabled by default (--prefix_cache_entries=0), non-greedy is untested, and the caveat is now in the header. Appropriate posture for landing as experimental.

Verdict

All three items from my prior review are addressed in this revision, and the open Copilot comments (lookup backoff, the const Session& lifetime, cache-hit metric distortion, the test-deadlock claim) are either stale against the current tree or resolved by @mergennachin's replies. I see no blocking bug in the exposed cloning path, the capture helper, or the prefix policy. With @metascroy's approval already in, this is ready from a correctness standpoint; the MLX numerical drift stays flagged as the thing to trace before enabling reuse by default or under non-greedy sampling.

--- • branch llm-prefix-cache-swa

[Copilot]

🟡 Changes recommended

The unresolved critical exception-handling and moderate capacity findings block approval.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (2)

backends/mlx/examples/llm/README.md:216

• The required cmp commands contradict the warning above that MLX cold/warm parity is not guaranteed: valid prefix reuse can change logits and generated text, including under greedy sampling, so this smoke test can fail despite correct cache behavior. Please describe these comparisons as optional diagnostics and make the request-report reuse checks the actual success criterion.

The generated text should agree across all four runs. In the cached run's
request report, generation 0 should have zero cached prompt tokens and generation 1
should reuse the prompt except for its final token, which is replayed for
fresh logits. Repeat with `--cache_kind batched-cell` to check the alternative

backends/mlx/examples/llm/run_llm_batched.cpp:367

• snapshot_slots is the maximum number of snapshots this example can ever retain (min(prompts.size(), prefix_cache_entries)), but this reservation adds the full configured entry count. With one prompt and, for example, 63 entries, the example asks the cache factory for 65 resident sessions even though it can create at most one retained snapshot; a bounded layout such as CellCache is rejected at its 64-sequence limit and the extra capacity is never usable. Reserve the actual possible retained-slot count here (or otherwise account for the prompt set) instead of FLAGS_prefix_cache_entries.

  const auto resident_sessions = static_cast<std::uint64_t>(prompts.size()) +
      FLAGS_prefix_cache_entries + snapshot_slots;

• Files reviewed: 19/19 changed files
• Comments generated: 1
• Review effort level: Lite