[mergennachin]

Reading a prompt is the slow first step of answering it: the model processes every token to build its internal state (prefill). With prefix caching on, the serving runtime saves that state as a snapshot, so a later request that starts with the same text can skip re-reading it.

Today the server saves a snapshot only for the first prompt of each session. Later turns aren't saved, so whenever the live session can't be reused, everything after the first prompt is read again. That happens when a turn is cancelled and retried, when another session (such as a sub-agent) continues the same conversation, or when an idle session is evicted and the client resends its history.

This PR adds an option to save every turn's prompt, not just the first:

Turn 1:  A B C          saved today
Turn 2:  A B C r D E    saved only with `every`

Turn 2 retried after a cancel, or reused by another session:
  today:          reuse A B C, re-read r D E
  with `every`:   reuse A B C r D, re-read only E

Use --prefix-checkpoint-prompt-end every to turn it on; initial, the default, keeps today's behavior. Snapshots go into the existing shared cache and count against the existing --prefix-cache-entries limit, so the memory reserved for a given setting doesn't change. It's off by default in both the generic and Muse-Glimmer servers, and it can't be combined with --prefix-checkpoint-interval.

Saving is best effort. The server skips prompts it has already saved, and it frees the temporary slot used to make a copy as soon as the copy exists, instead of holding it until the response finishes. Both behaviors also apply to the default mode.

An earlier version of this PR kept a separate cache for each session. Based on review feedback, it now uses the shared cache instead. That way requests without a session ID benefit too, and --prefix-cache-entries keeps a single meaning.

Review order: serving_runtime.h and prefix_cache.h, then serving_runtime.cpp. The rest is tests, launchers, and docs.

Tested with the serving and batching C++ suites (296 and 227 tests) and the launcher Python tests. The end-to-end Muse-Glimmer check from the earlier version hasn't been re-run yet.

Authored with OpenAI Codex and Claude Code.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 1 Unrelated Failure

As of commit 381baa0 with merge base 0fc43f6 ():

NEW FAILURES - The following jobs have failed:

• MLX / test-mlx-muse-glimmer-build / test-mlx-muse-glimmer-build (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 8
• Test CUDA Builds / unittest-cuda-muse-glimmer-batching / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 8

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Test CUDA Builds / test-muse-glimmer-cuda-e2e-kquant-17gb-dflash-image-batching / linux-job (gh) (trunk failure)
  ##[error]received SIGINT

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

[mergennachin]

@claude /pr-review

[claude]

Claude finished @mergennachin's task in 20m 7s —— View job

---

Summary: The new prefix_capture_prompt_end flag does capture continuation and replacement prompt ends, but I'm requesting changes. Turning it on also moves first-prompt capture onto the single checkpoint lane, which captures one prompt at a time, while the docs describe the flag as only adding captures. Two of the new tests also pass with the feature reverted.

Detailed review

PR Review: #24624

Summary

This PR adds an opt-in prefix_capture_prompt_end to ServingRuntimeConfig, with --prefix-capture-prompt-end in both launchers. It adds each request's prompt end to that request's prefill checkpoint list. Problems found:

• Turning the flag on means fewer first prompts are captured and fewer transient rows are available when requests overlap.
• Two new tests cannot detect a regression in the feature.
• Existing interface docs were not updated.

All 10 changed files were reviewed.

Code Quality

• serving_runtime.cpp:1233-1236 repeats the include_end logic of periodic_prefill_breakpoints (prefix_cache.h:60-62).

  • The block is only needed because the helper returns early when interval <= 0 (prefix_cache.h:51).
  • This is the helper's only caller outside tests. Let the helper honor include_end at interval 0 and pass opening || config_.prefix_capture_prompt_end.
  • The existing periodic_prefill_breakpoints(0, 8, 0) expectation still holds.
  • The combined checkpoint list could then be tested next to prefix_cache_test.cpp:431-447.
  • Fix this →

• The new user-facing text is inaccurate. In BATCHING_SERVING.md:146-158:

  • "Prefix caching retains the initial prompt by default" reads as if caching were on by default. It is opt-in (line 141).
  • "Benchmark … memory" is misleading. Physical rows are fixed at sessions + entries + resolved_prefix_captures(), and the flag only lowers the automatic row count.
  • "The capture lane" is what the code calls the checkpoint lane. In the code, "capture lanes" means the transient rows.
  • The prefix_checkpoint_policy sentence describes an option the Muse-Glimmer launcher cannot set.

  Elsewhere:

  • The help strings ("Retain each request's prompt end…") leave out the "best effort" caveat that the header comment includes.
  • module_worker_main.cpp:65 still says only "checkpoints require prefix_cache_entries > 0".
  • serve.py:46 no longer names the flag that caused the error.

  Fix this →

Testing

• prefix_reuse_test.cpp:1046-1051: the interval=4 case passes with the new code deleted.
  • The continuation starts at position 3, so periodic_prefill_breakpoints(3, 6, 4) already returns {4}.
  • That periodic snapshot alone satisfies minimum_reuse = 4 and prefilled <= 2.
  • Only the 0 and 64 cases catch a regression. So the "periodic checkpoint plus exact prompt end" combination advertised at BATCHING_SERVING.md:153-154 is never checked.
  • Fix: make the case deterministic by holding the tail until the periodic clone is collected, then assert reuse equals prompt.size() - 1. Otherwise, soften the doc.
  • Fix this →
• prefix_reuse_test.cpp:1103-1123: CancelledContinuationSkipsPromptEndCapture passes with the feature reverted.
  • Before this PR a continuation had no capture path, so clone_calls == before holds trivially.
  • The final {8, 9, 10} request on "other" opens a new session, and the default wrap path captures it anyway.
  • The cancel lands during prefill, before the end breakpoint can take the checkpoint lane. runner.cpp:1285-1290 finishes cancelled generations before checking breakpoints, and the lambda at serving_runtime.cpp:1240 rejects cancelled requests. So a cancel after the lane is held is never tested.
  • Fix: hold executor.cloning after the end breakpoint and cancel while the clone is pending. Then assert the lane is released and a later continuation is captured.
• prefix_reuse_test.cpp:1092-1101: ExplicitCheckpointPolicyOverridesPromptEndCapture cannot tell "the policy won" from "nothing worked".
  • The policy returns no checkpoints and neither request's result is checked, so the test also passes if both requests fail.
  • Fix: return a non-empty policy (for example, one checkpoint at 2), assert the only clone is at that position, and assert both requests succeeded.
• test_serve_launcher.py:170-191: only the (8, 0, 2) case depends on the flag. (1, 0, 1) and (8, 3, 3) give the same result with the flag off. A prompt-end column in the existing test_capture_capacity_validation_and_worker_override table (:194-236) would cover this without a second test.
• Missing coverage:
  • A named-session replacement that shares a prefix with history. This is the "branched" case from the PR description and goes through the clone_prefix path at serving_runtime.cpp:742. Only the replacement with no shared prefix is tested.
  • An image-prompt continuation, which BATCHING_SERVING.md:160 says takes the same path.

API Design

• serving_runtime.h:63-65: the new public field overlaps prefix_checkpoint_policy. Needs discussion.
  • The policy already receives start and opening and returns positions in (start, prompt.size()], so it can express "periodic checkpoints plus prompt end".
  • ready (serving_runtime.cpp:262-266) and resolved_prefix_captures() already treat a policy as checkpoint mode.
  • The new bool adds a third mode and a precedence rule.
  • The interval || prompt_end || policy check is now written out in ready, in generate() (serving_runtime.cpp:1197-1200) and in resolved_prefix_captures() (serving_runtime.h:92-94). Partial copies are in serve.py:42-54 and both worker validators.
  • Suggestion: have the launchers install a ready-made policy built from their CLI flags. If the field stays, give ServingRuntimeConfig one method for the three runtime checks.
  • Naming: elsewhere "capture" means the transient rows (max_prefix_captures, --prefix-capture-slots). This option goes through checkpoint admission, so prefix_checkpoint_prompt_end would match its sibling options.
• Existing interface docs now describe the behavior incorrectly.
  • serving_runtime.h:60-61 still says an interval of "zero captures only the initial prompt".
  • serving_runtime.h:80-81 presents the two-row cap as a checkpoint-mode rule without saying prompt-end mode is included.
  • examples/llm_server/spec/batching.md:34-37 still says "Zero (the default) captures only the initial prompt".
  • examples/llm_server/python/README.md:66-70 and examples/llm_server/spec/README.md:282-287 still describe "one pending prompt capture per admitted request" and "With checkpoints enabled…".
  • --prefix-capture-prompt-end is not documented for serve.py at all; only the Muse-Glimmer guide mentions it.

Performance

• serving_runtime.cpp:1197-1200: turning the flag on means fewer first prompts are captured when requests overlap.
  • Flag off: each request that opens a session takes the wrap branch (:1265-1281). It reserves its own transient row at submit through try_acquire_capture_lane(). Up to resolved_prefix_captures() of these can run at once.
  • Flag on: every text request takes the observer branch, so the wrap branch is never reached. All requests, including ones that open a session, compete for one checkpoint_lane_busy_ slot (:731-740). That slot stays held until the control thread collects the clone.
  • No retry: the runner moves past each breakpoint whether or not on_prefill captured anything (runner.cpp:1293-1304). A session that reaches its prompt end while another clone is pending is never captured.
  • Fewer rows: automatic transient rows drop from min(max_sessions, max_requests) to at most 2. Named-session replacements (clone_prefix, :743) compete for those same rows.
  • Why it matters: the first prompt of a new session was already captured without the flag, so the flag gains nothing there. Yet BATCHING_SERVING.md:146-148 says it will "also retain" later boundaries. This hits the prompt-end-only setup the PR recommends.
  • Comparison: interval mode makes the same trade-off, but its docs say collection is serialized (python/README.md:69-70).
  • Fix: keep requests that open a session on the wrap path when prompt-end is the only checkpoint option, since wrap already captures at the prompt end, and size rows to match. Otherwise, document the trade-off. Either way, add a test where several new sessions open at the same time.
  • Fix this →
• serving_runtime.cpp:1233-1236: exact repeats clone the full prompt and then throw the clone away.
  • Lookup caps a match at request.size() - 1 (prefix_cache.h:385).
  • So a repeated prompt, such as a retry or regenerate on an existing session, starts at end - 1, passes end > start, and queues a clone of all end rows.
  • insert_identity (prefix_cache.h:299-305) then finds the identical key, refreshes its recency and discards the new snapshot.
  • While that clone is pending it holds the single checkpoint lane, so another session's new boundary can be skipped.
  • The default path already does this for requests that open a session; the flag extends it to existing sessions.
  • Fix: skip the prompt-end breakpoint when the cache already holds an entry of exactly that length for the prompt.
  • Fix this →

Recommendation

Request Changes

• Compared with the default path, the flag captures fewer first prompts and leaves fewer transient rows when requests overlap, while the docs present it as purely additive. Either fix this or document it.
• The interval=4 case and the cancellation test should fail when the feature is reverted.
• The question of a new field versus a policy is worth settling before this experimental config gets another mode.

[mergennachin]

@claude /pr-review

[claude]

Claude finished @mergennachin's task in 20m 35s —— View job

---

Summary: This needs discussion before merge: in prompt-end-only mode each continuation and replacement holds a transient row until it finishes, the flag behaves differently alongside an interval or policy, the new capacity check is untested, and a helper change is unreachable.

Detailed review

PR Review: #24624

Summary

This PR adds opt-in prefix_capture_prompt_end / --prefix-capture-prompt-end, which also retains continuation and replacement prompt ends. Used alone, it captures through per-request capture lanes. Alongside an interval, it shares the checkpoint lane. Problems found:

• In prompt-end-only mode, each continuation and replacement keeps a transient row until it finishes, so explicit slot limits can force replacements into a cold prefill.
• The flag's behavior depends on whether an interval or a policy is also set, and validation accepts every combination.
• The test for the new capacity rule passes without that rule.
• The new interval-0 path in periodic_prefill_breakpoints has no caller.

All 15 changed files were reviewed.

Code Quality

• prefix_cache.h:51-61: the new interval-0 path in periodic_prefill_breakpoints is unreachable.
  • 2ea9667 made interval == 0 with include_end return {end}. The same commit moved prompt-end-only mode off the observer path.
  • The helper's only runtime caller (serving_runtime.cpp:1264) runs inside uses_checkpoint_observer() with no policy set, so the interval it passes is always positive.
  • The helper's comment (:44) does not describe the new interval-0 meaning. The (3, 6, 0, true) case at prefix_cache_test.cpp:443-444 tests behavior nothing uses.
  • Fix: restore the interval <= 0 early return and drop that case.
  • Fix this →

Testing

• prefix_reuse_test.cpp:1107-1114: PromptEndCaptureRequiresCapacityAndBoundsTransientRows never checks the new capacity rule.
  • prefix_checkpoint_interval is still 64 when start() runs, so the existing interval rule is what rejects the config.
  • Deleting !config.prefix_capture_prompt_end from serving_runtime.cpp:267-268 fails no C++ test. This is the only test that sets the flag with zero capacity.
  • Fix: reset the interval to 0 before start().
  • Fix this →

API Design

• prefix_capture_prompt_end behaves differently depending on the other two capture options. Needs discussion.
  • Alone: each request takes its own capture lane (serving_runtime.cpp:1303-1308).
  • With an interval: the end goes through the single checkpoint lane (:1276-1280, :747-756), often right after the same request's last periodic point. When the lane is busy, observer drops the breakpoint with no retry (prefix_cache.h:133-135). The end is lost whenever control has not yet collected an earlier clone, whether that clone came from the same request or another one.
  • Test: the interval=4 case only passes deterministically because it holds the prefill tail and uses an open_session_async("s") round trip as a control-thread barrier (prefix_reuse_test.cpp:1056-1065). Nothing in the test explains the barrier.
  • With prefix_checkpoint_policy: the header says the policy overrides the flag, but valid_config_ (serving_runtime.cpp:259-269) still accepts the combination, and the flag is then ignored (:1248).
  • Docs: all four docs now need a paragraph about these combinations. BATCHING_SERVING.md:244-245 advises against using the flag with an interval.
  • Options: reject both combinations at validation, or capture the prompt end through a per-request lane in every mode and size rows to match.
• examples/llm_server/spec/README.md:254 still describes the batching.md link as covering "creation-only caching".
  • The PR already rewrote the same phrase at :390.
  • batching.md now documents continuation and replacement capture.
  • Fix this →

Performance

• serving_runtime.cpp:1303-1310: in prompt-end-only mode, every continuation and replacement holds a capture lane for its whole generation.
  • wrap clones at the first output token (prefix_cache.h:95-101). The lane is released only by finish_capture(), which finalize() calls (serving_runtime.cpp:1095-1103, :1347). Without the flag, continuations hold no lane.
  • Named-session replacements need the same lanes. When none is free, clone_prefix returns Unavailable (:758-761), and prepare_text falls back to a cold reset and a full prefill (:990-992, :1007-1009) instead of using prefix_clone.
  • Trigger: --prefix-capture-slots set below the number of requests decoding at once. BATCHING_SERVING.md:346-349 tells users to reduce slot counts when the reservation does not fit. Automatic sizing gives one lane per concurrent request, so the default configuration is not affected.
  • Same cause: the snapshot enters the cache only at finalize, so a branch or replay submitted while the source turn is still decoding cannot reuse that prompt end.
  • Fix: collect the snapshot and release the lane once the clone resolves, as the observer path does through pending_checkpoint_. Otherwise, document the row lifetime in the max_prefix_captures comment (serving_runtime.h:78-83) and the launcher docs.
  • Fix this →

Recommendation

Needs Discussion

• Decide whether to shorten the capture-lane lifetime or document it.
• Decide whether the flag may be combined with an interval or a policy.
• The test fix and the helper revert are small and do not depend on either decision.

[mergennachin]

@claude /pr-review

[claude]

Claude finished @mergennachin's task in 33m 18s —— View job

---

Request Changes: turning per-session prompt-end checkpoints on by default for Muse-Glimmer quadruples the default executor rows (4 → 16), which breaks existing 65536-cell CUDA exports and command lines. Separately, --prefix-cache-entries now has two meanings, and the close/reset cleanup test never reaches the discard path it is named for.

Detailed review

PR Review: #24624

Summary

The PR gives each named session in ServingRuntime its own prompt-end PrefixCache, and enables that mode by default for Muse-Glimmer. The snapshot is captured at the first output token and collected before decode finishes. Problems:

• The new default breaks existing Muse-Glimmer exports and command lines.
• prefix_cache_capacity now carries a second meaning.
• Two of the new tests don't exercise the behavior they're named for.

Code Quality

• Dead test gate and redundant assignments (extension/llm/serving/test/prefix_reuse_test.cpp:170-171):

  • The new prefill_tail / gated_prefill_position gate is dead. gated_prefill_position is never assigned, so the check at 261-262 never fires and the release() at 451 does nothing.
  • Lines 1310-1312 assign prefix_cache_capacity = 8 twice. Lines 1369-1371 assign 8 and then overwrite it with GetParam().

  Fix this →

• Unrelated doc edits in examples/llm_server/python/README.md and examples/llm_server/spec/README.md, none of them about prompt-end checkpoints. Revert them or move them to a separate PR:

  • Glimmer/Muse mentions are removed.
  • The QWEN_HF_DIR/GEMMA_HF_DIR/MUSE_GLIMMER_HF_DIR list is replaced with a pointer to a test file.
  • The launcher-status sentences are rewritten.
  • One edit is wrong: spec/README.md:142 changes "The generic and Muse batching launchers accept --api-dialect" to "The batching launcher accepts…". Both launchers still accept it (python/serve.py:146, muse-glimmer/serving/serve_batching.py:121).

  Fix this →

• Second queue pattern (minor, extension/llm/serving/serving_runtime.cpp:1205-1218):

  • The class's other queues are containers. This adds an intrusive list instead (RequestState::next_capture, pending_captures_, pending_captures_tail_).
  • The comment justifies it "so the engine callback needs no additional allocation". But the engine's settle callback already allocates through schedule() → dispatch_queue_.push_back (1467-1473, 1568).
  • Use a container like dispatch_queue_. If the real reason is that on_capture in PromptCapture::wrap must not throw, the comment should say that.

  Fix this →

Testing

• CloseAndResetClearPendingAndRetainedSnapshots never reaches collect(false) (extension/llm/serving/test/prefix_reuse_test.cpp:1429-1451):

  • The engine queues the capture before the clone runs, and the control loop serves pending_captures_ before commands (serving_runtime.cpp:1790-1799).
  • So collect_capture computes retain while "s" still exists, then blocks in collect(true) until executor.cloning.release(). Close/reset is processed only after that.
  • The final reused_prompt_tokens == 0 assertion also holds without the clear() calls (serving_runtime.cpp:411-412, 459-461). The next open of "s" gets a fresh per-session cache, and the dropped cache gives back its row once the request releases its last reference, before the open_count poll finishes.
  • The test therefore stays green if the retain check (1188-1195) or the clear() calls are deleted.
  • It needs an ordering in which collect_capture runs with retain == false, plus an open_count() check before the source request completes.

  Fix this →

• EarlyPromptCaptureTest case 0 doesn't test early release (prefix_reuse_test.cpp:1276-1292):

  • It probes a different session ("probe"), which with per-session caches can never reuse the checkpoint of "source".
  • So its assertions (reused_prompt_tokens == 0, "new") hold whether or not the capture lane was released before decode finished.
  • The "Reuse" in the instantiation name ReuseAndReplacementAndRefusedCapture (1301-1304) is left over from shared-cache semantics.
  • Drop case 0, or replace it with a case whose result depends on the lane being free while "source" is paused.

  Fix this →

API Design

• prefix_cache_capacity / --prefix-cache-entries now has two meanings (extension/llm/serving/serving_runtime.h:58-63). It is the shared limit normally and the per-session limit in prompt-end mode, where the shared cache is built with capacity 0 (serving_runtime.cpp:281). As a result:

  • The shared cache, which handles anonymous and cross-session reuse, can't be used together with per-session checkpoints.
  • The runtime never reads kDefaultSessionCheckpoints (serving_runtime.h:40), so a C++ config that only sets prefix_checkpoint_prompt_end = true is invalid.
  • The default of 2 is resolved separately in four places: python/serve.py:157-159, serve_batching.py:132-135, and two is_default checks that overwrite the global FLAGS_prefix_cache_entries (module_worker_main.cpp:55-58, batching_bootstrap.cpp:70-74). Neither native branch is tested.
  • --help for the Muse-Glimmer worker shows the declared default of 0 next to help text saying "(default: 2)" (batching_bootstrap.cpp:35-38).
  • A dedicated field and flag defaulting to kDefaultSessionCheckpoints would remove all four resolutions and leave --prefix-cache-entries with one meaning.

  Fix this →

• --prefix-checkpoint-prompt-end has different syntax in the two launchers. python/serve.py:129 defines it as store_true, so =false is a usage error there. serve_batching.py:89-96 accepts [=true|false]. Use one definition in both.

  Fix this →

• PromptCapture::collect(bool retain = true) is a mode flag on a public ET_EXPERIMENTAL API (extension/llm/batching/prefix_cache.h:174-178). collect(false) collects nothing: it waits for the clone and drops it. A separately named method would be clearer. The private process(Command&, bool keep_checkpoints = false) (serving_runtime.cpp:385, one caller at 1112) uses the same pattern.

  Fix this →

Backward Compatibility

• The Muse-Glimmer default raises physical rows from 4 to 16. Prompt-end checkpoints are now on by default (serve_batching.py:89-96, batching_bootstrap.cpp:47-50). The default config now needs 4 sessions + 4×2 snapshots + 4 capture slots (BATCHING_SERVING.md:252). After upgrading:

  • CUDA exports with the default 65536 cells (BATCHING_SERVING.md:333) reject any context above 4095 tokens (cuda_executor.cpp:486). The documented 8192-token setup now needs --max-cells 131073 (BATCHING_SERVING.md:328, 364-367), which doubles the KV pool.
  • When the native worker runs with max_session_tokens = 0, its automatic context drops from 16383 to 4095 tokens without any error (cuda_executor.cpp:374-376).
  • On MLX the cache is sized rows × session_capacity (mlx_executor.cpp:217), so the default budget quadruples.
  • The CUDA session limit (CellCache::kMaxSeqs = 64 rows, cuda_executor.cpp:516-522) falls from 64 sessions to 20 at the default 4 in-flight requests, or to 16 when in-flight requests equal sessions.
  • Command lines that used to work are now rejected: any --prefix-checkpoint-interval N, and --prefix-cache-entries 0.
  • An explicit --prefix-cache-entries N silently becomes N per session and disables the shared cache. The --prefix-cache-entries 32 setup that this PR removes from the Muse-Glimmer README grows from 40 to 136 rows, past the CUDA limit.
  • Keep the mode opt-in, as in the generic launcher. If on-by-default is intended, an explicit --prefix-checkpoint-interval N or --prefix-cache-entries 0 should turn it off instead of failing.

  Fix this →

Recommendation

Request Changes

What blocks approval:

• The Muse-Glimmer default breaks existing exports and command lines, with no migration path.
• The overloaded prefix_cache_capacity / --prefix-cache-entries should be split before more launchers depend on it.
• The cleanup test should fail when the retain check or the clear() calls are removed.

---

• Gather PR context and diff
• Run three independent reviewers
• Consolidate findings
• Fact-check findings (5 reviewer claims dropped as invalid)
• Post review

Review only: no code was changed, and nothing was built or run.

[kiymetakdemir]

Prefix caching configuration feels like getting complicated to me. Enabling --prefix-checkpoint-prompt-end saves every session request’s prompt end instead of only the initial prompt, and changes --prefix-cache-entries 8 from 8 checkpoints total to 8 per session.

What do you think about separating checkpoint ownership from capture policy?

Ownership: shared across sessions, or private per session.
Capture policy: initial prompt end, every request’s prompt end, or interval checkpoints.

These can be expressed separately, for example:
--checkpoint-scope shared|session
--checkpoint-prompt-end initial|every
--checkpoint-interval 512 # optional
--checkpoint-entries N

The entry limit would apply to the selected scope. We can still reject the unsupported combinations.

However, separating ownership from capture policy might not worth redesigning if we don’t plan to support more combinations. For example, do we expect to support other capture policies with per-session ownership in the future, or is per-session caching intended to be limited to prompt-end checkpoints? Could be a follow up work as well.

[mergennachin]

Agreed, it got too complicated. Rather than splitting, I'm removing per-session ownership from this PR. If we need per-session ownership later, --prefix-cache-scope shared|session can be added without breaking anything, with the entry limit applying to the selected scope as you suggested. We'll have three options:

--prefix-cache-entries N: total shared checkpoints; 0 turns caching off
--prefix-checkpoint-prompt-end initial|every (new, default initial): every also saves each later request's prompt end
--prefix-checkpoint-interval N Same as main; works with initial only

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

[mergennachin]

@claude /pr-review

[claude]

Claude finished @mergennachin's task in 29m 18s —— View job

---

Needs discussion: with --prefix-checkpoint-prompt-end every, one active conversation's per-turn snapshots can evict other sessions' prompt ends from the shared cache. The other findings are minor, and the Muse-Glimmer end-to-end check still needs a re-run on the final head.

Detailed review

PR Review: #24624

Summary

The PR adds PromptEndCapture { Initial, Every } (--prefix-checkpoint-prompt-end) so the serving runtime can keep the prompt end of every request in the shared PrefixCache, not only requests that open a session. It also releases the capture lane once the clone settles, instead of when decoding finishes. I found no correctness blockers. The main issue is how every interacts with the global LRU. There is also one test that no longer exercises the path its name describes, a missing assertion in the new cancellation test, and two small code-quality points.

Code Quality

• The run loop chooses how to release a capture from global config, not from the request. extension/llm/serving/serving_runtime.cpp:1755-1762
  • What happens: when config_.uses_checkpoint_observer() is true, it takes the observer release (collect_capture, fetch_sub, clear checkpoint_lane_busy_). Otherwise it calls finish_capture.
  • Why it's fragile: this is only correct because the observer branch at 1326-1327 always wins over the prompt-end branch at 1384 when an observer is configured. If that condition changes, a prompt-end capture would go through the observer release with capture_lane still set. Then finalize → finish_capture would decrement capture_lanes_used_ a second time, and the transient-row limit would no longer hold.
  • Fix: dispatch on capture->text->capture_lane instead. Fix this →
• serve.py forwards the new flag in a different form from the options next to it.
  • What happens: examples/llm_server/python/serve.py:233 passes a single --prefix_checkpoint_prompt_end=every token. The adjacent options, such as --prefix_checkpoint_interval at 225-228, are passed as separate [flag, value] tokens, and tests/test_serve_launcher.py:157 checks them that way.
  • Fix: pass the flag and its value as two tokens, and update the expected command in tests/test_serve_launcher.py:171 and 238. Fix this →

Testing

• BusyCaptureLaneAllowsWarmLookupsWithoutLosingCaptureOwnership no longer tests a busy capture lane. extension/llm/serving/test/prefix_reuse_test.cpp:2926-2971
  • What happens: the warm requests send {1, 2, 3} again, and the seed request already put that prompt in the cache. The !prefix_cache_.contains_identity(...) check (serving_runtime.cpp:1389) therefore skips them before try_acquire_capture_lane() (1390) runs. The "no new capture" assertions at 2938 and 2970 would pass even if the lane were free.
  • Fix: use warm prompts that aren't already cached, sent while the owner still holds its lane. Alternatively, rename the test to describe what it now checks. Fix this →
• The new cancellation test doesn't check what happens to the snapshot. CancellationDuringPromptEndCloneReleasesAdmission (prefix_reuse_test.cpp:1165-1189)
  • What happens: the test checks that the capture lane is released, but not what happens to the snapshot from the cancelled request's clone. collect_capture (serving_runtime.cpp:1152-1165) doesn't check for cancellation, so finalize → finish_capture still adds {1, 2, 100, 3} to the cache.
  • Fix: send that prompt again under a new session key and expect reused_prompt_tokens == 3. If dropping the snapshot is intended, assert that the lookup misses instead. Fix this →

Performance

• With every, one active conversation can fill the shared cache.

  • Cause: in every mode, extension/llm/serving/serving_runtime.cpp:1384-1390 captures the prompt end of every turn of a live session. insert_identity (extension/llm/batching/prefix_cache.h:323-330) only merges exact duplicates, so each turn adds another entry that contains the previous one. Live continuations never look up or touch() the cache, so the LRU order is just insertion order.

  • Example: --prefix-cache-entries 4 --prefix-checkpoint-prompt-end every. Session A goes idle after one turn, and session B runs four turns. B's four entries push A's prompt end out of the cache. If A's live session is later evicted for being idle and its history is resent, it gets no reuse. Under initial it would have.

  • Waste: on executors that can clone at any earlier position, B's three older entries are never useful, because find_identity always prefers the longest match.

  • Docs: examples/models/muse-glimmer/BATCHING_SERVING.md:237-241 warns that "new boundaries can evict earlier prompts" but gives no sizing guidance.

  • Options:

    • Document sizing: entries ≈ sessions to keep × captured turns per session.
    • Or, when the executor can clone at arbitrary positions, have a new capture replace the entry it extends. Replacement isn't safe on executors that "can clone only at saved boundaries" (extension/llm/batching/runner.h:212), because there the older entries are the only exact restore points.

    Fix this →

Recommendation

Needs Discussion

• Decide how every handles a session's older entries: document the sizing, or replace the older entries where the executor allows it.
• Re-run the end-to-end check: the PR description says the Muse-Glimmer check hasn't been re-run since the earlier version. Re-run it on the final head with every enabled before merging.
• The rest is minor: the testing and code-quality items don't block approval.

Review steps

• Gather PR context and diff (9 commits, 15 files, +764/−97 against origin/main)
• Run three independent reviewers (all 15 files covered)
• Consolidate findings (14 candidates; 4 reported by at least two reviewers)
• Fact-check findings (5 kept, 8 dropped, and the end-to-end re-run moved to the Recommendation)
• Post review

This was a review only: no code was changed, and nothing was built or run.

---