[caikun-pjlab]

Motivation

Chunked prefill for hybrid attention models can create multiple intermediate state checkpoints and Mooncake Store submissions within one input turn.

For append-only agent conversations, subsequent turns typically reuse the latest prefix. Saving fewer intermediate checkpoints can reduce state copies and Store writes, at the cost of fewer reusable prefixes for edited or branching conversations.

Modification

• Add prefill_save_last_only / --prefill-save-last-only, disabled by default.
• When enabled, select the last reusable boundary aligned to prefill_save_alignment for each input turn. Preserve the required input suffix and move the boundary backward when necessary to avoid splitting multimodal inputs, including shifted MTP inputs.
• Save the local prefix checkpoint and submit the state snapshot plus all unsaved full-attention KV blocks through that boundary to Mooncake Store, when the respective caches are enabled. This boundary can occur before the final prefill chunk.
• Skip intermediate checkpoints. If no valid boundary exists or a cache hit already reaches it, do not create another prefill checkpoint.
• Keep normal chunk budgets before the selected boundary and update the scheduler’s chunk-count estimate accordingly.
• Change the default prefill_save_alignment from 8192 to 1024.
• Move prefix-caching documentation from the English and Chinese speculative-decoding guides into dedicated prefix_caching.md pages.

The new policy applies to hybrid attention models. Full-attention-only models ignore it.

Compatibility

The last-only policy is opt-in. The alignment default changes independently; deployments requiring the previous alignment should explicitly set: --prefill-save-alignment 8192

Earlier turns’ checkpoints remain subject to normal eviction. Snapshot-slot exhaustion retains the existing skip behavior.

MTP configurations that rewind cache hits by full blocks may require additional KV blocks beyond the saved state boundary for remote reuse. This limitation is documented.

Validation

Precision

Dataset	vLLM (%)	LMDeploy (%)	Difference (pp)
GPQA-Diamond — average of 8 runs	91.35	91.16	-0.19
AIME 2025 — average of 32 runs	96.56	96.46	-0.10
GSM8K	97.65	97.5	-0.38
MMMU-Pro	77.69	78.61	+0.92

Performance

Added test coverage for configuration propagation, chunk planning, multimodal and MTP boundaries, coordinated local/Store saves, and reuse across turns.

Benchmarked GLM-5.3-Flash FP8 on 8 × H200 with TP8 and MTP enabled:

• Client concurrency: 1.
• Mooncake OFF/ON × baseline/last-only.
• Five service launches per configuration, with two measured samples per test point per launch.
• Every launch runs 14 warmup requests before measurement.
• All 560 measured requests passed input/output length, output consistency, and cache-hit checks.
• No Triton cache files were added or updated during measurement.

Mooncake ON results below. TTFT changes are relative to baseline; negative values indicate lower latency.

Scenario	Store write reduction	First/single-turn TTFT change	Subsequent-turn TTFT change
Single turn, 8K + 64 input	0.00%	-1.77%	—
Single turn, 16K + 64 input	42.43%	-0.69%	—
Single turn, 32K + 64 input	63.64%	+0.32%	—
Single turn, 64K + 64 input	74.25%	+0.18%	—
Five turns, starting at 32K + 64, adding 4K per turn	33.07%	-0.05%	-1.51%
Five turns, starting at 32K + 64, adding 16K per turn	49.50%	-0.89%	-1.19%

Single-turn output is fixed at 16 tokens; multi-turn output is fixed at 512 tokens. Store bytes are successful put payloads summed across all TP ranks, with multi-turn reductions calculated over the complete conversation.

The measured benefit is reduced checkpoint and Store write volume. All paired TTFT difference 95% intervals cross zero, so this experiment does not establish a consistent single-concurrency TTFT improvement. Full-attention KV write volume remains unchanged.

The benchmark covers local prefix reuse and Store writes. Remote restoration, branching conversations, multimodal performance, and cache-pressure scenarios are outside its scope.

[RunningLeon]

LGTM

[yidingcheng0206]

[P2] Account for block-rewind requirements when selecting the only checkpoint
This boundary reserves input-shift/lookahead tokens but ignores seq.prefix_cache.recompute_overlap.recompute_blocks. For a 2049-token prompt with alignment=1024, block_size=64 and recompute_blocks=1, last-only saves state at 2048. An identical local request cannot reuse it because _find_recompute_overlap_end() requires a matching block through 2112. The earlier checkpoint at 1024 remains reusable with last-only disabled, but is no longer saved here. Please select a boundary that satisfies the actual restore requirements and add a local cold/warm regression test with recompute_blocks=1 and token_lookahead=0.

[caikun-pjlab]

With #5061 merged, all MTP methods use token_lookahead=1 and recompute_blocks=0 for local prefix caching. I verified #5049 on top of that change: the 2049-token warm request now reuses the checkpoint at 2048, so the reported local MTP case no longer reproduces.

[qescccczmr]

LGTM