feat(pytorch): add last-only prefill checkpoint saving for hybrid attention

docs/en/advance/prefix_caching.md

@@ -0,0 +1,103 @@
+# Prefix caching (PyTorch)
+
+This guide covers local prefix caching in the PyTorch backend and the hybrid
+state-checkpoint policies shared with Mooncake Store. Local caching is enabled
+with `PytorchEngineConfig.enable_prefix_caching=True`; Mooncake Store must be
+configured separately through `kv_transfer_config`.
+For TurboMind, see its [prefix caching settings](../inference/turbomind_config.md#prefix-caching-switch).
+
+## Prefill checkpoints for hybrid attention models
+
+Hybrid attention models need both full-attention KV blocks and the
+linear-attention state at a reusable prefix boundary.
+
+- `prefill_save_alignment` (`--prefill-save-alignment`, default 1024) controls
+  prefill state-checkpoint boundaries for both local prefix caching and Mooncake
+  Store. It must be a positive multiple of the cache block size.
+  With alignment 512 and a token budget of 8192,
+  an uncached 800-token text prompt uses two chunks: 512 and 288 tokens,
+  including when both caches are enabled. Full-attention-only models ignore this
+  option. Alignment does not require a checkpoint at every multiple.
+- `prefix_cache_state_budget` reserves extra checkpoint slots and uses additional
+  memory. Zero adds no reserved slots; idle runtime slots may still be reused.
+
+### Save only the last checkpoint
+
+`prefill_save_last_only=True` (`--prefill-save-last-only`, disabled by default)
+saves only the last reusable boundary aligned to `prefill_save_alignment` in
+each input turn. The boundary leaves at least one input token and moves back
+if needed to avoid splitting multimodal inputs, including the shifted MTP
+input. Earlier chunks use their normal token budget; the chunk reaching this
+boundary saves the local state checkpoint and sends all unsaved full-attention
+KV blocks through that boundary, together with its state snapshot, to Mooncake
+Store. Each cache must be enabled separately.
+
+For example, with alignment explicitly set to 8192 and 20,000 input tokens,
+the save is at 16,384; if an image spans \[16,000, 18,000), it moves back to 8192.
+A 32,768-token input saves at 24,576, leaving a suffix for prefix reuse.
+No save is made if no positive safe boundary exists or a cache hit already
+reaches it. Earlier turns' checkpoints remain subject to normal eviction,
+and unavailable snapshot slots retain the existing skip behavior.
+
+This policy reduces checkpoint copies and Store submissions for append-only
+agent conversations, at the cost of fewer intermediate prefixes for edited
+or branching conversations. Full-attention-only models ignore it.
+
+MTP configurations that rewind cache hits by full blocks still require KV
+beyond the saved state boundary for remote reuse. This policy does not save
+that suffix, so the latest state may be unusable remotely without those blocks
+already in Store. GLM MTP with local prefix caching uses token lookahead
+instead of block rewind and does not require that extra Store margin.
+
+## GLM-5.3-Flash
+
+GLM-5.3-Flash can reuse prefix caches with ordinary autoregressive decoding or
+[`deepseek_mtp`](./spec_decoding.md#deepseek-mtp).
+Enable `PytorchEngineConfig.enable_prefix_caching` in either mode.
+The following text-only example uses TP4; choose a TP size and cache budget that
+fit your hardware. Set `speculative_config=None` to disable MTP without disabling
+prefix caching.
+
+```python
+from lmdeploy import GenerationConfig, PytorchEngineConfig, pipeline
+from lmdeploy.messages import SpeculativeConfig
+
+if __name__ == '__main__':
+    engine = PytorchEngineConfig(
+        tp=4,
+        max_batch_size=4,
+        enable_prefix_caching=True,
+        prefix_cache_state_budget=8,
+        max_prefill_token_num=512,
+        language_model_only=True,
+    )
+    spec = SpeculativeConfig(method='deepseek_mtp', num_speculative_tokens=3)
+    with pipeline('/path/to/GLM-5.3-Flash', backend_config=engine,
+                  speculative_config=spec) as pipe:
+        notes = 'A rectangle has four sides and four right angles. ' * 80
+        prompt = notes + '\nHow many sides does a rectangle have?'
+        generation = GenerationConfig(max_new_tokens=128, do_sample=False)
+        cold = pipe(prompt, gen_config=generation)
+        warm = pipe(prompt, gen_config=generation)
+        print(cold.cached_tokens, warm.cached_tokens)
+```
+
+### Configuration limits
+
+- Cache blocks contain 64 tokens. If setting `num_gpu_blocks` manually, count
+  these blocks rather than KPool entries.
+- With MTP, keep `prefix_cache_decode_state_interval=0`. Without MTP, positive
+  values must be multiples of 64.
+- PD migration is not supported.
+
+The native FP8 CUDA path also supports DP/EP serving, for example `--tp 2 --dp 2 --ep 4`. Enable text-prefill PCG with `--piecewise-cudagraph-max-tokens 512`
+and `--max-prefill-token-num 512`. Startup warmup prepares these plans; skipping
+warmup leaves unprepared prefills eager. Vision-bearing chunks also fall back to
+eager execution. MTP draft prefill remains eager, while supported decode steps
+use the ordinary CUDA graph path. These choices are independent of prefix caching.
+
+Inspect `Response.cached_tokens` to verify actual reuse. For aggregate local
+and Store hit rates, see [cache hit metrics](./metrics.md).
+Changing prefill, PCG padding, batch shapes or parallel topology can change
+FP8 generation trajectories; compare task accuracy and MTP acceptance as well
+as cache hits rather than assuming bitwise-identical outputs or a guaranteed speedup.

docs/en/advance/spec_decoding.md

@@ -145,66 +145,6 @@ When a DFlash block size is provided, it overrides
 `--speculative-num-draft-tokens` by setting the number of newly proposed
 tokens to `block_size - 1`.
 
-## GLM-5.3-Flash prefix caching
-
-GLM-5.3-Flash can reuse prefix caches with ordinary autoregressive decoding or
-`deepseek_mtp`. Enable `PytorchEngineConfig.enable_prefix_caching` in either mode.
-The following text-only example uses TP4; choose a TP size and cache budget that
-fit your hardware. Set `speculative_config=None` to disable MTP without disabling
-prefix caching.
-
-```python
-from lmdeploy import GenerationConfig, PytorchEngineConfig, pipeline
-from lmdeploy.messages import SpeculativeConfig
-
-if __name__ == '__main__':
-    engine = PytorchEngineConfig(
-        tp=4,
-        max_batch_size=4,
-        enable_prefix_caching=True,
-        prefix_cache_state_budget=8,
-        max_prefill_token_num=512,
-        language_model_only=True,
-    )
-    spec = SpeculativeConfig(method='deepseek_mtp', num_speculative_tokens=3)
-    with pipeline('/path/to/GLM-5.3-Flash', backend_config=engine,
-                  speculative_config=spec) as pipe:
-        notes = 'A rectangle has four sides and four right angles. ' * 80
-        prompt = notes + '\nHow many sides does a rectangle have?'
-        generation = GenerationConfig(max_new_tokens=128, do_sample=False)
-        cold = pipe(prompt, gen_config=generation)
-        warm = pipe(prompt, gen_config=generation)
-        print(cold.cached_tokens, warm.cached_tokens)
-```
-
-### Configuration limits
-
-- Cache blocks contain 64 tokens. If setting `num_gpu_blocks` manually, count
-  these blocks rather than KPool entries.
-- `prefill_save_alignment` (`--prefill-save-alignment`, default 8192) controls
-  prefill state-checkpoint boundaries for both local prefix caching and Mooncake
-  Store. It must be a positive multiple of the cache block size.
-  With alignment 512 and a token budget of 8192,
-  an uncached 800-token text prompt uses two chunks: 512 and 288 tokens,
-  including when both caches are enabled. Full-attention-only models ignore this
-  option. Alignment does not require a checkpoint at every multiple.
-- `prefix_cache_state_budget` reserves extra checkpoint slots and uses additional
-  memory. Zero adds no reserved slots; idle runtime slots may still be reused.
-- With MTP, keep `prefix_cache_decode_state_interval=0`. Without MTP, positive
-  values must be multiples of 64.
-- PD migration is not supported.
-
-The native FP8 CUDA path also supports DP/EP serving, for example `--tp 2 --dp 2 --ep 4`. Enable text-prefill PCG with `--piecewise-cudagraph-max-tokens 512`
-and `--max-prefill-token-num 512`. Startup warmup prepares these plans; skipping
-warmup leaves unprepared prefills eager. Vision-bearing chunks also fall back to
-eager execution. MTP draft prefill remains eager, while supported decode steps
-use the ordinary CUDA graph path. These choices are independent of prefix caching.
-
-Inspect `Response.cached_tokens` to verify actual reuse. Changing prefill,
-PCG padding, batch shapes or parallel topology can change FP8 generation trajectories; compare task accuracy and
-MTP acceptance as well as cache hits rather than assuming bitwise-identical
-outputs or a guaranteed speedup.
-
 ## Guided Decoding with Speculative Decoding
 
 Speculative decoding (MTP) can be combined with [structured output](./structed_output.md) so that the draft tokens proposed by the spec model also respect the grammar constraints (e.g. JSON schema, regex). This significantly improves the acceptance rate compared to running spec decoding without grammar masks.

docs/en/index.rst

@@ -117,6 +117,7 @@ Documentation
    advance/pytorch_profiling.md
    advance/metrics.md
    advance/context_parallel.md
+   advance/prefix_caching.md
    advance/spec_decoding.md
    advance/update_weights.md
    advance/checkpoint_engine.md

docs/zh_cn/advance/prefix_caching.md

@@ -0,0 +1,86 @@
+# 前缀缓存（PyTorch）
+
+本文介绍 PyTorch 后端的本地前缀缓存，以及与 Mooncake Store 共用的 hybrid
+state checkpoint 保存策略。通过 `PytorchEngineConfig.enable_prefix_caching=True`
+启用本地缓存；Mooncake Store 需通过 `kv_transfer_config` 单独配置。
+TurboMind 后端请参阅其[前缀缓存配置](../inference/turbomind_config.md#前缀缓存开关)。
+
+## Hybrid attention 模型的 prefill checkpoint
+
+Hybrid attention 模型复用前缀时，需要同时具备 full-attention KV blocks
+和该前缀边界上的 linear-attention state。
+
+- `prefill_save_alignment`（CLI 参数 `--prefill-save-alignment`，默认 1024）统一控制
+  本地 prefix caching 和 Mooncake Store 的 prefill 状态快照分块边界，取值必须是缓存块大小的正整数倍。设置为 512、
+  token 预算为 8192 时，首次处理 800-token 纯文本输入会切成 512 和 288 两段，
+  同时开启两种缓存也一样。仅含 full attention 的模型忽略此参数；对齐不代表每到一个倍数都保存快照。
+- `prefix_cache_state_budget` 为 checkpoint 预留额外状态槽，会增加显存占用。
+  设为零表示不额外预留，仍可复用空闲的运行时状态槽。
+
+### 仅保存最后一个 checkpoint
+
+`prefill_save_last_only=True`（CLI 参数 `--prefill-save-last-only`，默认关闭）
+每轮输入仅保存最后一个对齐 `prefill_save_alignment` 的可复用边界。
+该边界至少保留一个输入 token，遇到多模态 span 时向前回退，并考虑 MTP 的输入偏移。
+此前的 chunk 按正常 token 预算推进；到达目标边界的 chunk 同时保存本地 state checkpoint，
+并向 Mooncake Store 发送截至该边界尚未保存的全部 full-attention KV blocks 和对应 state snapshot。
+两种缓存仍需分别启用。
+
+例如显式将对齐设为 8192，输入长度 20,000 时保存在 16,384；
+若图像占据 \[16,000, 18,000)，则回退到 8192。输入长度恰为 32,768 时保存在 24,576，
+留出后缀用于 prefix 复用。没有正数安全边界或缓存命中已到达该边界时，本轮不新增保存。
+旧轮次的 checkpoint 仍按原有规则淘汰，snapshot 槽不足时仍按现有策略跳过。
+
+该选项可减少追加式 agent 多轮对话的 checkpoint 拷贝和 Store 提交次数，
+代价是编辑历史或分叉对话时可复用的中间前缀更少。仅含 full attention 的模型忽略此参数。
+
+若 MTP 配置仍采用整块回退，远端复用还需要 state 边界之后的 KV。
+此策略不保存这段后缀，因此 Store 中尚无这些 blocks 时，最新 state 可能无法远端命中。
+GLM MTP 开启本地 prefix caching 时使用 token lookahead，不采用整块回退，无需这段额外 Store KV。
+
+## GLM-5.3-Flash
+
+GLM-5.3-Flash 的普通自回归解码和 [`deepseek_mtp`](./spec_decoding.md#deepseek-mtp) 都可以启用
+`PytorchEngineConfig.enable_prefix_caching`。下面是 TP4 的纯文本示例，
+请根据硬件选择 TP 和缓存预算。将 `speculative_config` 设为 `None` 即可关闭
+MTP，同时保留 prefix caching。
+
+```python
+from lmdeploy import GenerationConfig, PytorchEngineConfig, pipeline
+from lmdeploy.messages import SpeculativeConfig
+
+if __name__ == '__main__':
+    engine = PytorchEngineConfig(
+        tp=4,
+        max_batch_size=4,
+        enable_prefix_caching=True,
+        prefix_cache_state_budget=8,
+        max_prefill_token_num=512,
+        language_model_only=True,
+    )
+    spec = SpeculativeConfig(method='deepseek_mtp', num_speculative_tokens=3)
+    with pipeline('/path/to/GLM-5.3-Flash', backend_config=engine,
+                  speculative_config=spec) as pipe:
+        notes = 'A rectangle has four sides and four right angles. ' * 80
+        prompt = notes + '\nHow many sides does a rectangle have?'
+        generation = GenerationConfig(max_new_tokens=128, do_sample=False)
+        cold = pipe(prompt, gen_config=generation)
+        warm = pipe(prompt, gen_config=generation)
+        print(cold.cached_tokens, warm.cached_tokens)
+```
+
+### 配置限制
+
+- 缓存块大小为 64 tokens。手动设置 `num_gpu_blocks` 时按此计数，而不是按 KPool 条目计数。
+- MTP 下保持 `prefix_cache_decode_state_interval=0`；非 MTP 的正值必须是 64 的倍数。
+- 暂不支持 PD 迁移。
+
+原生 FP8 CUDA 路径也支持 DP/EP 服务，例如 `--tp 2 --dp 2 --ep 4`。
+使用 `--piecewise-cudagraph-max-tokens 512 --max-prefill-token-num 512`
+启用文本 prefill PCG。启动 warmup 负责准备 plan；禁用 warmup 时未准备的
+prefill 走 eager，包含视觉输入的 chunk 也回退到 eager。MTP draft prefill
+仍走 eager，受支持的 decode 使用普通 CUDA graph。这些选项与 prefix caching 独立。
+
+用 `Response.cached_tokens` 检查实际命中，本地和 Store 的累计命中率见[缓存命中指标](./metrics.md)。
+Prefill 分块、PCG padding、batch 形状或并行拓扑变化可能改变 FP8 的生成轨迹；
+应同时检查任务 accuracy、MTP accept rate 和缓存命中，不能默认逐 token 完全相同或一定获得性能收益。

docs/zh_cn/advance/spec_decoding.md

@@ -142,59 +142,6 @@ Qwen/Qwen3.5-35B-A3B \
 设置 DFlash block size 后，它会覆盖 `--speculative-num-draft-tokens`，
 并将新提出的 token 数设置为 `block_size - 1`。
 
-## GLM-5.3-Flash 前缀缓存
-
-GLM-5.3-Flash 的普通自回归解码和 `deepseek_mtp` 都可以启用
-`PytorchEngineConfig.enable_prefix_caching`。下面是 TP4 的纯文本示例，
-请根据硬件选择 TP 和缓存预算。将 `speculative_config` 设为 `None` 即可关闭
-MTP，同时保留 prefix caching。
-
-```python
-from lmdeploy import GenerationConfig, PytorchEngineConfig, pipeline
-from lmdeploy.messages import SpeculativeConfig
-
-if __name__ == '__main__':
-    engine = PytorchEngineConfig(
-        tp=4,
-        max_batch_size=4,
-        enable_prefix_caching=True,
-        prefix_cache_state_budget=8,
-        max_prefill_token_num=512,
-        language_model_only=True,
-    )
-    spec = SpeculativeConfig(method='deepseek_mtp', num_speculative_tokens=3)
-    with pipeline('/path/to/GLM-5.3-Flash', backend_config=engine,
-                  speculative_config=spec) as pipe:
-        notes = 'A rectangle has four sides and four right angles. ' * 80
-        prompt = notes + '\nHow many sides does a rectangle have?'
-        generation = GenerationConfig(max_new_tokens=128, do_sample=False)
-        cold = pipe(prompt, gen_config=generation)
-        warm = pipe(prompt, gen_config=generation)
-        print(cold.cached_tokens, warm.cached_tokens)
-```
-
-### 配置限制
-
-- 缓存块大小为 64 tokens。手动设置 `num_gpu_blocks` 时按此计数，而不是按 KPool 条目计数。
-- `prefill_save_alignment`（CLI 参数 `--prefill-save-alignment`，默认 8192）统一控制
-  本地 prefix caching 和 Mooncake Store 的 prefill 状态快照分块边界，取值必须是缓存块大小的正整数倍。设置为 512、
-  token 预算为 8192 时，首次处理 800-token 纯文本输入会切成 512 和 288 两段，
-  同时开启两种缓存也一样。仅含 full attention 的模型忽略此参数；对齐不代表每到一个倍数都保存快照。
-- `prefix_cache_state_budget` 为 checkpoint 预留额外状态槽，会增加显存占用。
-  设为零表示不额外预留，仍可复用空闲的运行时状态槽。
-- MTP 下保持 `prefix_cache_decode_state_interval=0`；非 MTP 的正值必须是 64 的倍数。
-- 暂不支持 PD 迁移。
-
-原生 FP8 CUDA 路径也支持 DP/EP 服务，例如 `--tp 2 --dp 2 --ep 4`。
-使用 `--piecewise-cudagraph-max-tokens 512 --max-prefill-token-num 512`
-启用文本 prefill PCG。启动 warmup 负责准备 plan；禁用 warmup 时未准备的
-prefill 走 eager，包含视觉输入的 chunk 也回退到 eager。MTP draft prefill
-仍走 eager，受支持的 decode 使用普通 CUDA graph。这些选项与 prefix caching 独立。
-
-用 `Response.cached_tokens` 检查实际命中。Prefill 分块、PCG padding、batch 形状或并行拓扑变化可能
-改变 FP8 的生成轨迹；应同时检查任务 accuracy、MTP accept rate 和缓存命中，
-不能默认逐 token 完全相同或一定获得性能收益。
-
 ## 投机解码与结构化输出
 
 投机解码（MTP）可以与[结构化输出](./structed_output.md)结合使用，使草稿模型提出的 token 也遵循语法约束（如 JSON Schema、正则表达式），从而显著提高接受率。

docs/zh_cn/index.rst

@@ -111,6 +111,7 @@ LMDeploy 工具箱提供以下核心功能：
    advance/pytorch_profiling.md
    advance/metrics.md
    advance/context_parallel.md
+   advance/prefix_caching.md
    advance/spec_decoding.md
    advance/update_weights.md
    advance/checkpoint_engine.md

lmdeploy/cli/cli.py

@@ -57,6 +57,7 @@ def add_parser_chat():
         ArgumentHelper.prefix_cache_state_budget(pt_group)
         ArgumentHelper.prefix_cache_decode_state_interval(pt_group)
         ArgumentHelper.prefill_save_alignment(pt_group)
+        ArgumentHelper.prefill_save_last_only(pt_group)
         ArgumentHelper.mooncake_state_save_slots(pt_group)
         # common engine args
         dtype_act = ArgumentHelper.dtype(pt_group)

lmdeploy/cli/serve.py

@@ -121,6 +121,7 @@ def add_parser_api_server():
         ArgumentHelper.prefix_cache_state_budget(pt_group)
         ArgumentHelper.prefix_cache_decode_state_interval(pt_group)
         ArgumentHelper.prefill_save_alignment(pt_group)
+        ArgumentHelper.prefill_save_last_only(pt_group)
         ArgumentHelper.mooncake_state_save_slots(pt_group)
 
         # common engine args
@@ -267,6 +268,7 @@ def api_server(args):
                 prefix_cache_state_budget=args.prefix_cache_state_budget,
                 prefix_cache_decode_state_interval=args.prefix_cache_decode_state_interval,
                 prefill_save_alignment=args.prefill_save_alignment,
+                prefill_save_last_only=args.prefill_save_last_only,
                 mooncake_state_save_slots=args.mooncake_state_save_slots,
                 device_type=args.device,
                 quant_policy=args.quant_policy,