[Exp][Hybrid KV Cache] Evaluate input-end vs decode-end checkpoint retention

[qianlihuang]

Summary

This experimental PR evaluates semantic endpoint admission for recurrent-state prefix caching in mamba_cache_mode=align.

The production baseline is the current main-branch sparse-retention policy with:

VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0

Under this policy, main retains:

• the literal prompt or replay endpoint, referred to here as P_end;
• reactively discovered Marconi-style shared-prefix junctions.

Relative to the current sparse-retention baseline, this experiment evaluates two proactive endpoint candidates:

• Input-end checkpoint (U_end): the deepest materializable recurrent state at or before the end of the stable externally supplied conversation history, before generation-only prompt suffixes are appended.
• Decode-end checkpoint (D_end): the deepest materializable recurrent state at or before the end of the committed generated token prefix.

The three logical boundaries are:

stable externally supplied history
                                  ^ U_end
generation-only prompt suffix
                              ^ P_end
committed generated output
                           ^ D_end

A later request continues to use the existing chained token-prefix hashes to select the longest correct checkpoint.

Phase 1 retains both U_end and D_end. Its purpose is to measure their production hit distribution, reusable-token benefit, recurrent-state storage cost, and eviction externality before introducing a policy that attempts to select only one.

Existing reactive shared-prefix retention remains enabled in every experiment variant.

Motivation

Hybrid attention models combine paged attention KV cache with large recurrent states.

The existing fine-grained hybrid prefix-cache infrastructure already provides the mechanisms required to register and match recurrent states at prefix_match_unit boundaries. The remaining question is which materializable states should be admitted under sparse retention.

The current sparse main-branch policy retains the literal prompt or replay endpoint:

stable conversation history
                           ^ U_end
assistant generation header
reasoning-generation prefix
                           ^ P_end

For reasoning models, P_end may contain generation-only tokens such as:

• an assistant generation header;
• a reasoning-mode control token;
• a <think> prefix;
• other template content used to initiate generation.

These tokens may not appear in the next canonical rendering after the generated assistant response becomes conversation history.

For example:

Request N:

[user] question
               ^ U_end
[assistant generation header]
<think>
       ^ P_end

The next request may instead render:

Request N+1:

[user] question
               ^ U_end remains reachable
[assistant history header]
[visible assistant answer]
[user] follow-up

In this case, the token chain represented by P_end is not a prefix of the next request.

The attention cache may still contain blocks before the divergence, but a hybrid cache can only resume at a position where the required attention and recurrent cache groups are jointly available.

If no earlier recurrent checkpoint has already been materialized, the effective joint prefix hit on the first continuation may therefore fall to zero.

Main already mitigates repeated branch misses through Marconi-style reactive shared-prefix checkpointing. Once a shared branch has been observed, main can materialize and retain the corresponding recurrent state.

This experiment evaluates whether proactive request-semantic endpoints can avoid the first miss before such a branch has necessarily been observed.

Checkpoint definitions

Literal prompt end (P_end)

P_end is the existing sparse-retention endpoint represented by the end of the literal tokenized prompt or replay sequence.

It may include generation-only chat-template content after the stable externally supplied conversation history.

Examples include:

[assistant generation header]
[reasoning mode prefix]
<think>

P_end is reusable only when its complete token chain remains an exact prefix of a later request.

Input end (U_end)

U_end is the logical token position immediately after the stable externally supplied conversation history and before generation-only prompt suffixes.

Typical external history includes user messages and tool results, but the definition is not tied to a specific role. It may also include system, developer, assistant-history, or application-inserted messages that are already part of the canonical request history.

The materialized U_end checkpoint is the deepest valid prefix_match_unit boundary at or before that logical token position.

U_end is intentionally conservative. A longer generation-prefix segment may also remain stable for some chat templates, but this experiment does not attempt to infer that longer boundary.

Decode end (D_end)

D_end is the logical end of the committed generated token prefix.

The materialized D_end checkpoint is the deepest valid prefix_match_unit boundary at or before that position.

It excludes:

• rejected speculative tokens;
• uncommitted draft tokens;
• tokens rolled back during generation;
• terminal or control tokens that are not part of the committed reusable prefix.

D_end is reusable only when the generated sequence and its later serialized chat representation reproduce the same token prefix.

Textual equivalence is insufficient: chat-template rendering, role tokens, parser output, escaping, and structured-output serialization must collectively preserve the token chain.

Examples

Tool-call continuation: D_end is the expected longest hit

During a tool loop, the generated assistant output may be included in the next request before the tool result.

Request N:

[user] Find the weather in Shanghai.
                                   ^ U_end
[assistant reasoning]
[assistant tool_call: get_weather("Shanghai")]
                                                  ^ D_end

The next request appends the tool result:

Request N+1:

[user] Find the weather in Shanghai.
[assistant reasoning]
[assistant tool_call: get_weather("Shanghai")]
                                                  ^ D_end hit
[tool] 31°C, partly cloudy

If parsing and chat rendering reproduce the generated prefix token-for-token, D_end is the longest reusable checkpoint.

If the serialized representation changes, chained token hashes reject D_end and lookup falls back to an earlier matching checkpoint such as U_end or an existing shared-prefix checkpoint.

Reasoning response: U_end is the expected fallback

For a normal reasoning response, the model may execute hidden reasoning that is omitted from later conversation history.

Request N execution:

[user] Explain why the sky is blue.
                                    ^ U_end
[assistant hidden reasoning]
[assistant] Rayleigh scattering causes shorter wavelengths...
                                                               ^ D_end

The next request may contain only the visible assistant response:

Request N+1 rendering:

[user] Explain why the sky is blue.
[assistant] Rayleigh scattering causes shorter wavelengths...
[user] Does this also explain sunsets?

D_end cannot be reused because its recurrent state depends on hidden reasoning tokens that are absent from the new token prefix.

The chained token hash rejects D_end, and lookup falls back to U_end:

[user] Explain why the sky is blue.
                                    ^ U_end hit

The visible assistant answer and the new user message are then recomputed.

Reasoning disabled or preserved

If reasoning is disabled, or if generated reasoning remains part of the exact serialized history, the complete committed output may remain a prefix of the next request:

Request N:

[user] question
               ^ U_end
[assistant] answer
                  ^ D_end

Request N+1:

[user] question
[assistant] answer
                  ^ D_end hit
[user] follow-up

Phase 1 does not attempt to predict which behavior applies. Both checkpoints are retained, and token hashes select the longest correct match.

Relationship to existing main-branch behavior

The relevant baseline is main with sparse endpoint retention:

VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0

The baseline retains:

• P_end, the literal prompt or replay endpoint;
• reactively discovered shared-prefix junctions.

Historically, hybrid align mode also retained the latest materialized recurrent state, conceptually corresponding to Marconi's last-state checkpoint for chat continuation. PR #37898 added reactive shared-prefix admission on top of that pre-existing last-state behavior.

Selective sparse retention later changed retention_interval=0 to retain only the latest prompt replay boundary. PR #47782 preserved reactive shared-prefix junctions under that sparse policy, but did not restore the decode-side last state.

Therefore, D_end is not a new checkpoint category relative to Marconi. In this experiment it explicitly restores Marconi-style last-state retention under the current sparse-retention baseline. U_end is the additional semantic candidate introduced to handle continuations where the literal prompt or generated execution suffix does not remain an exact prefix.

This PR does not disable or replace reactive shared-prefix retention.

The experimental variable is the proactive endpoint policy:

main_sparse:
    P_end

input_end:
    U_end

decode_end:
    D_end

both:
    U_end + D_end

All variants continue to use existing Marconi-style shared-prefix detection and retention.

Default dense retention is not a primary comparison target because retaining recurrent states at many aligned positions can be prohibitively expensive for models with large recurrent state.

Relationship to Marconi

Main already implements Marconi-style reactive shared-prefix checkpointing for hybrid prefix caching.

That mechanism discovers a useful branch after observing divergent requests and materializes the recurrent state required to make the shared prefix jointly reusable.

The relationship differs between the two evaluated endpoints:

• D_end corresponds to Marconi's last-state policy. The contribution here is to restore that checkpoint explicitly under selective sparse retention, where the current retention_interval=0 baseline otherwise retains the prompt replay boundary instead.
• U_end is complementary to Marconi. It proactively retains an earlier stable-history boundary for cases where generation-only prompt tokens, hidden reasoning, or output reserialization make P_end and D_end unreachable from the next canonical request.

Together, the experiment evaluates whether restoring last-state retention and adding a stable-history fallback can avoid the first continuation miss before a shared branch has necessarily been observed:

Marconi-style retention:

retain the latest state
    -> accelerate exact chat continuation

observe a divergent shared prefix
    -> materialize the shared junction
    -> improve later branch reuse


Semantic endpoint retention:

restore D_end under sparse retention
    -> preserve Marconi-style exact continuation

identify stable externally supplied history
    -> retain U_end as a fallback
    -> handle non-canonical generation suffixes

U_end targets generation-only or non-canonical prompt suffixes that can make P_end unreachable.

D_end targets exact continuation workloads in which committed model output round-trips token-for-token into the next request.

This PR does not attempt to implement Marconi's complete admission, scoring, or eviction policy. It evaluates endpoint retention on top of the existing reactive shared-prefix mechanism.

Deriving U_end

Phase 1 does not derive U_end through counterfactual prompt comparison.

The logical U_end token offset must be supplied by a structured request-rendering path that knows when stable externally supplied history ends and generation-only prompt content begins.

Conceptually:

render stable supplied messages
    -> record U_end
append generation-only template content
    -> reach P_end

This experiment intentionally avoids:

• chat-template mutation probes;
• secondary counterfactual rendering;
• character longest-common-prefix inference;
• token longest-common-prefix inference;
• arbitrary semantic checkpoint inference inside opaque prompts.

Opaque rendering paths that cannot expose the boundary between stable history and generation-only suffixes are outside the initial scope.

Phase 1

Phase 1 retains both U_end and D_end.

The purpose is to establish an empirical upper bound for reuse from these two endpoint candidates and measure the cost of retaining the second recurrent state.

A both run can measure which candidate is selected, but it is insufficient on its own because retaining an additional state changes:

• cache pressure;
• checkpoint residency;
• eviction behavior;
• the number of sessions that fit in the cache.

Separate policy runs are therefore required.

Experiment variants

Every policy retains existing reactive shared-prefix checkpoints.

Policy	Proactive endpoint retention	Reactive shared-prefix retention
main_sparse	P_end	enabled
input_end	U_end	enabled
decode_end	D_end	enabled
both	U_end and D_end	enabled

The initial implementation in this PR uses both.

The full experiment should evaluate the variants under two comparison modes.

Equal total cache capacity

Every variant receives the same total recurrent-state memory capacity.

This measures realistic production behavior, including the additional eviction pressure caused by retaining two endpoints.

Equal admitted checkpoint-byte budget

Every variant receives the same admitted recurrent-state byte budget.

Under this comparison, both may retain fewer sessions or may need to evict one candidate earlier.

This isolates the incremental utility of the second endpoint from the advantage of consuming more state memory.

Questions

This experiment is intended to answer:

1. What is the effective joint prefix hit on the first continuation after a checkpoint is produced?
2. How often does the next request reuse D_end?
3. How often does D_end reuse come from tool-call continuation?
4. When D_end does not match, how often does U_end provide a useful fallback?
5. How often does U_end avoid a first-continuation miss caused by an unreachable P_end?
6. After a shared branch has been observed and reactively materialized, how much additional benefit does U_end retain?
7. How many additional prompt tokens are reused by each policy?
8. How much prefill work is avoided by each policy?
9. How much recurrent-state memory does each endpoint consume?
10. Does the second checkpoint increase cache eviction enough to offset its reuse benefit?
11. What is the effect on TTFT, throughput, and cache residency under realistic concurrency?
12. How does behavior differ between first reuse and steady-state reuse?
13. How often do multiple logical candidates map to the same materialized checkpoint?

Metrics

The experiment should record:

Checkpoint selection

• selected logical checkpoint:
  • prompt_end;
  • input_end;
  • decode_end;
  • reactive shared-prefix checkpoint;
• selected materialized token boundary;
• producer request age:
  • immediately preceding request;
  • two requests earlier;
  • older;
• checkpoint consumability in the current cache domain.

Prefix reuse

• attention-group matched token length;
• recurrent-group matched token length;
• effective joint matched token length;
• exact chained-hash matched token length;
• additional matched tokens relative to main_sparse;
• prompt tokens recomputed after the hit;
• estimated avoided prefill work;
• first-continuation hit rate;
• steady-state hit rate;
• reactive branch-learning misses.

Checkpoint cost

• retained recurrent-state bytes;
• checkpoint creation count;
• checkpoint hit count;
• checkpoint lifetime;
• checkpoint snapshot or materialization bytes;
• checkpoint materialization time;
• forced scheduler split count;
• additional scheduler or model-execution steps caused by checkpoint placement;
• duplicate logical candidates mapping to one physical checkpoint;
• eviction count and reason;
• displaced checkpoint count;
• cache occupancy.

End-to-end performance

• TTFT;
• request throughput;
• input-token throughput;
• cache residency;
• prefix-cache hits and misses.

Useful derived metrics include:

decode_end_hit_rate
input_end_fallback_rate
first_continuation_effective_hit_rate
saved_prompt_tokens_per_retained_state_byte
saved_prefill_work_per_retained_state_byte
incremental_hit_rate_of_second_checkpoint
incremental_saved_work_of_second_checkpoint

Correctness

Checkpoint placement does not determine correctness.

The existing chained token-prefix hashes continue to validate the exact token prefix represented by each cached recurrent state.

A checkpoint that includes hidden, rewritten, discarded, rejected, or otherwise non-canonical tokens simply fails to match a later request.

The request then falls back to the next longest checkpoint that is jointly available across all required cache groups.

No recurrent state may be reused beyond the validated token prefix it represents.

Relationship to fine-grained hybrid prefix caching

This PR does not introduce a new partial-prefix cache representation or lookup mechanism.

It builds on the existing hybrid prefix-cache infrastructure:

• physical cache-block size remains independent from prefix_match_unit;
• chained hashes identify complete token prefixes;
• recurrent states can be retained at materializable boundaries;
• attention and recurrent cache groups converge on the same logical hit length;
• partial hits continue through the existing copy-on-write path.

The experimental variable is which request-boundary states are proactively retained.

Relationship to chat-aware checkpointing

This PR does not yet implement a chat-aware single-endpoint selection policy.

Phase 1 retains both candidates so that their actual hit distribution can be measured.

A later phase may use this data to choose one checkpoint:

generated output is expected to remain an exact token prefix:
    retain D_end

generated output contains non-persistent reasoning or is expected to be rewritten:
    retain U_end

A parser-aware extension may retain the end of the longest decode prefix that is serialized verbatim into conversation history.

More general prompt-internal or decode-internal checkpoint placement should only be considered if production measurements show that U_end and D_end leave substantial reusable work uncaptured.

Planned phases

Phase 1: retain both endpoints

Retain:

• U_end;
• D_end.

Continue retaining existing reactive shared-prefix checkpoints.

Measure:

• first-continuation and steady-state hit behavior;
• selected checkpoint distribution;
• saved work;
• storage cost;
• eviction externality;
• end-to-end performance.

Phase 2: select one endpoint

Evaluate whether a chat-aware or parser-aware policy can choose between U_end and D_end without materially reducing reuse.

Possible initial policy:

• select D_end for active tool-call continuation or output expected to round-trip token-for-token;
• select U_end when generated reasoning or generation-only suffixes are not expected to remain in canonical history.

The policy should be evaluated under the same checkpoint-byte budget as main_sparse.

Phase 3: finer checkpoint placement

Only if endpoint selection is insufficient, evaluate more precise locations such as:

• the end of the longest generated prefix preserved verbatim by the parser;
• boundaries before non-persistent control-token suffixes;
• explicitly supplied application-level semantic boundaries.

Canonical replay and arbitrary chat-template probing remain separate design options.

PD-disaggregated deployments

D_end is useful to a later prefill request only when both requests can access the same cache domain.

This may be provided by:

• colocated prefill and decode;
• sticky session routing;
• decode-to-prefill cache transfer;
• publication into a shared cache store accessible by prefill workers.

In a strictly disaggregated deployment with no decode-to-prefill publication path, D_end cannot benefit the next prefill request.

Such deployments should disable decode-end retention or report the checkpoint as non-consumable.

U_end may still be useful if it is produced and retained in the prefill cache domain.

The experiment should distinguish:

checkpoint retained
checkpoint hash-matched
checkpoint consumable by the target worker
checkpoint selected as the effective joint hit