From d8b20a153504da8b64eb72d05b5b2b7863989c02 Mon Sep 17 00:00:00 2001 From: caikun-pjlab Date: Wed, 21 Oct 2026 10:36:15 +0800 Subject: [PATCH] feat(kvcache): support only save last checkpoint to reduce TTFT and prefix checkpoint --- docs/en/advance/prefix_caching.md | 103 ++++++++++++++++++ docs/en/advance/spec_decoding.md | 60 ---------- docs/en/index.rst | 1 + docs/zh_cn/advance/prefix_caching.md | 86 +++++++++++++++ docs/zh_cn/advance/spec_decoding.md | 53 --------- docs/zh_cn/index.rst | 1 + lmdeploy/cli/cli.py | 1 + lmdeploy/cli/serve.py | 2 + lmdeploy/cli/utils.py | 12 +- lmdeploy/messages.py | 9 +- lmdeploy/pytorch/config.py | 3 +- lmdeploy/pytorch/engine/config_builder.py | 1 + .../kv_connector/mooncake/store/scheduler.py | 3 + lmdeploy/pytorch/long_context.py | 37 ++++++- lmdeploy/pytorch/messages.py | 1 + .../paging/block_trie/checkpoint_lifecycle.py | 3 + lmdeploy/pytorch/paging/prefill_scheduler.py | 15 ++- lmdeploy/pytorch/paging/scheduler.py | 1 + tests/pytorch/engine/test_executor_base.py | 10 +- .../test_mooncake_store_scheduler.py | 93 ++++++++++++++-- .../pytorch/paging/test_prefill_scheduler.py | 51 ++++++++- 21 files changed, 411 insertions(+), 135 deletions(-) create mode 100644 docs/en/advance/prefix_caching.md create mode 100644 docs/zh_cn/advance/prefix_caching.md diff --git a/docs/en/advance/prefix_caching.md b/docs/en/advance/prefix_caching.md new file mode 100644 index 0000000000..6855730e92 --- /dev/null +++ b/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. diff --git a/docs/en/advance/spec_decoding.md b/docs/en/advance/spec_decoding.md index ac6566df89..c10ad1463e 100644 --- a/docs/en/advance/spec_decoding.md +++ b/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. diff --git a/docs/en/index.rst b/docs/en/index.rst index a735815154..2837c8d561 100644 --- a/docs/en/index.rst +++ b/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 diff --git a/docs/zh_cn/advance/prefix_caching.md b/docs/zh_cn/advance/prefix_caching.md new file mode 100644 index 0000000000..f50f1a9654 --- /dev/null +++ b/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 完全相同或一定获得性能收益。 diff --git a/docs/zh_cn/advance/spec_decoding.md b/docs/zh_cn/advance/spec_decoding.md index d6f243d39b..649b5be5f6 100644 --- a/docs/zh_cn/advance/spec_decoding.md +++ b/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、正则表达式），从而显著提高接受率。 diff --git a/docs/zh_cn/index.rst b/docs/zh_cn/index.rst index bb40504721..fa7c6049e0 100644 --- a/docs/zh_cn/index.rst +++ b/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 diff --git a/lmdeploy/cli/cli.py b/lmdeploy/cli/cli.py index bf2d1e9a22..5de6046dfb 100644 --- a/lmdeploy/cli/cli.py +++ b/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) diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py index 6af5441658..1ae993413b 100644 --- a/lmdeploy/cli/serve.py +++ b/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, diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index c026811191..7bdc824484 100644 --- a/lmdeploy/cli/utils.py +++ b/lmdeploy/cli/utils.py @@ -643,13 +643,23 @@ def prefill_save_alignment(parser): """Add the shared prefill state-checkpoint alignment.""" return parser.add_argument('--prefill-save-alignment', type=int, - default=8192, + default=1024, help='Token alignment for linear-attention prefill state checkpoints, shared by ' 'local prefix caching and Mooncake Store. Must be a positive multiple of cache ' 'block size; intermediate aligned positions may be skipped. ' 'Ignored by full-attention-only models. ' 'Only used by the PyTorch engine.') + @staticmethod + def prefill_save_last_only(parser): + """Add the final prefill checkpoint policy.""" + return parser.add_argument('--prefill-save-last-only', + action='store_true', + help='Save local prefix checkpoints and Mooncake caches only at the last safe ' + 'prefill-save-alignment boundary before the end of each prefill. ' + 'Leaves at least one input token for prefix reuse. ' + 'Only used by PyTorch models with linear-attention state caches.') + @staticmethod def mooncake_state_save_slots(parser): """Add the bounded hybrid Mooncake snapshot capacity.""" diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index 3f22316ed7..b4a7e67db2 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -503,9 +503,13 @@ class PytorchEngineConfig: must be multiples of the cache block size. prefill_save_alignment: Token alignment for prefill state checkpoints, shared by local prefix caching and Mooncake Store. Defaults to - 8192 and must be a positive multiple of block_size. Only applies + 1024 and must be a positive multiple of block_size. Only applies to models with linear-attention state caches; full-attention-only models ignore it. Intermediate aligned positions may be skipped. + prefill_save_last_only: Save local prefix checkpoints and Mooncake + caches only at the last safe aligned boundary before the end of + each prefill. Leaves at least one input token for prefix reuse. + Defaults to False. Only applies to linear-attention state caches. mooncake_state_save_slots: Number of temporary hybrid state snapshots reserved for asynchronous Mooncake saves. Defaults to 8. A full snapshot pool skips the save without blocking inference. @@ -576,7 +580,8 @@ class PytorchEngineConfig: enable_prefix_caching: bool = False prefix_cache_state_budget: int = 0 prefix_cache_decode_state_interval: int = 0 - prefill_save_alignment: int = 8192 + prefill_save_alignment: int = 1024 + prefill_save_last_only: bool = False mooncake_state_save_slots: int = 8 device_type: str = 'cuda' eager_mode: bool = False diff --git a/lmdeploy/pytorch/config.py b/lmdeploy/pytorch/config.py index 75b2540b12..b88f2959a4 100644 --- a/lmdeploy/pytorch/config.py +++ b/lmdeploy/pytorch/config.py @@ -171,7 +171,8 @@ class CacheConfig: num_state_caches: int = None prefix_cache_state_budget: int = 0 prefix_cache_decode_state_interval: int = 0 - prefill_save_alignment: int = 8192 + prefill_save_alignment: int = 1024 + prefill_save_last_only: bool = False mooncake_state_save_slots: int = 8 # Model-owned checkpoint validity constraint, independent of prefill chunking. prefix_cache_block_aligned: bool = False diff --git a/lmdeploy/pytorch/engine/config_builder.py b/lmdeploy/pytorch/engine/config_builder.py index 6b998e96c3..83bd9abbc2 100644 --- a/lmdeploy/pytorch/engine/config_builder.py +++ b/lmdeploy/pytorch/engine/config_builder.py @@ -86,6 +86,7 @@ def build_cache_config(engine_config: PytorchEngineConfig): prefix_cache_state_budget=engine_config.prefix_cache_state_budget, prefix_cache_decode_state_interval=engine_config.prefix_cache_decode_state_interval, prefill_save_alignment=engine_config.prefill_save_alignment, + prefill_save_last_only=engine_config.prefill_save_last_only, mooncake_state_save_slots=engine_config.mooncake_state_save_slots, quant_policy=engine_config.quant_policy, device_type=engine_config.device_type, diff --git a/lmdeploy/pytorch/kv_connector/mooncake/store/scheduler.py b/lmdeploy/pytorch/kv_connector/mooncake/store/scheduler.py index ffd0e58526..29f283bd4c 100644 --- a/lmdeploy/pytorch/kv_connector/mooncake/store/scheduler.py +++ b/lmdeploy/pytorch/kv_connector/mooncake/store/scheduler.py @@ -15,6 +15,7 @@ KVLoadResult, RequestId, ) +from lmdeploy.pytorch.long_context import get_last_prefill_save_step from lmdeploy.pytorch.paging.state_manager import StateAllocator from lmdeploy.utils import get_logger @@ -311,6 +312,8 @@ def _build_save_requests( if full_blocks <= first_block: continue if self._is_hybrid: + if self._cache_config.prefill_save_last_only and token_len != get_last_prefill_save_step(request): + continue if (token_len % self._cache_config.prefill_save_alignment != 0 or not request.is_prefix_cache_boundary_safe(token_len)): continue diff --git a/lmdeploy/pytorch/long_context.py b/lmdeploy/pytorch/long_context.py index c39b008937..9988068af8 100644 --- a/lmdeploy/pytorch/long_context.py +++ b/lmdeploy/pytorch/long_context.py @@ -3,6 +3,7 @@ from collections import defaultdict from dataclasses import dataclass +from itertools import chain from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -95,10 +96,34 @@ def get_long_context_chunk_capacity(seq: 'SchedulerSequence', max_prefill_token_ return max(max_prefill_token_num, max_span) +def get_last_prefill_save_step(seq: 'SchedulerSequence') -> int | None: + """Return the last reusable aligned state boundary for the whole input + turn. + + Use the fixed input end and full multimodal history so chunk progress and cache hits cannot move the boundary. The + shifted draft input must stay with its multimodal span, even when local prefix caching is disabled. + """ + alignment = seq._seq_meta.prefill_save_alignment + if not alignment: + return None + input_shift = max(seq.prefill_input_shift, seq._seq_meta.prefix_cache_token_lookahead) + step = (seq.input_end_pos - max(1, input_shift)) // alignment * alignment + if step <= 0: + return None + + spans = [(max(0, span.start - input_shift), span.end) + for span in chain(seq.prefix_cache.multimodal_spans, seq.history_embeddings.embeddings)] + # Descending starts let each rewind visit any earlier overlapping span next. + for start, end in sorted(spans, reverse=True): + if start < step < end: + step = start // alignment * alignment + return step or None + + def limit_long_context_checkpoint(seq: 'SchedulerSequence', limit: int) -> int: """Cut the final prefill suffix after a reusable aligned checkpoint.""" alignment = seq._seq_meta.prefill_save_alignment - if alignment and seq._seq_meta.enable_prefix_caching: + if alignment and seq._seq_meta.enable_prefix_caching and not seq._seq_meta.prefill_save_last_only: # Matching leaves at least one input for logits, so an identical prompt # can reuse the preceding aligned state instead of its final state. step = ((seq.input_end_pos - 1) // alignment) * alignment @@ -118,8 +143,8 @@ def plan_long_context_chunk(seq: 'SchedulerSequence', include_multimodals: bool = True) -> LongContextChunkPlan: """Plan a model-safe chunk, optionally stopping at its last save boundary. - Alignment is a constraint, not a sampling interval: earlier aligned positions in this chunk are skipped. If no safe - aligned position advances the request, use the ordinary chunk and keep only its running state. + With last-only saving, ordinary chunks need no alignment; only the turn's final checkpoint cuts a chunk. Otherwise, + alignment is a constraint, not a sampling interval: earlier aligned positions in this chunk are skipped. """ chunk_size = min(seq.num_token_ids, chunk_limit) start = seq.num_history_ids @@ -140,8 +165,12 @@ def plan_long_context_chunk(seq: 'SchedulerSequence', break save_alignment = seq._seq_meta.prefill_save_alignment + if seq._seq_meta.prefill_save_last_only: + save_end = get_last_prefill_save_step(seq) + if save_end is not None and start < save_end < end: + end = save_end # Embedding-only inputs do not yet have a Store content identity. - if save_alignment and len(seq.history_embeddings) == 0: + elif save_alignment and len(seq.history_embeddings) == 0: save_end = end // save_alignment * save_alignment for span_start, span_end in reversed(chunk_spans): if span_start < save_end < span_end: diff --git a/lmdeploy/pytorch/messages.py b/lmdeploy/pytorch/messages.py index 4203265b24..595584885c 100644 --- a/lmdeploy/pytorch/messages.py +++ b/lmdeploy/pytorch/messages.py @@ -215,6 +215,7 @@ class SequenceMeta: enable_mooncake_store: bool = False # Zero disables state-checkpoint chunking, including for full-attention-only models. prefill_save_alignment: int = 0 + prefill_save_last_only: bool = False prefix_cache_checkpoint_block_size: int = 0 prefix_cache_token_lookahead: int = 0 diff --git a/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py b/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py index 124ccdbe34..5f3ec09066 100644 --- a/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py +++ b/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py @@ -19,6 +19,7 @@ import numpy as np +from lmdeploy.pytorch.long_context import get_last_prefill_save_step from lmdeploy.pytorch.messages import SchedulerSequence from lmdeploy.utils import get_logger @@ -166,6 +167,8 @@ def reserve_save(self, seq: SchedulerSequence, step: int = None, is_decode: bool if step is None: step = seq.num_valid_ids + if not is_decode and seq._seq_meta.prefill_save_last_only and step != get_last_prefill_save_step(seq): + return -1 lookahead = seq._seq_meta.prefix_cache_token_lookahead if lookahead and (is_decode or step + lookahead > seq.input_end_pos or seq.logprob_start_pos >= 0): return -1 diff --git a/lmdeploy/pytorch/paging/prefill_scheduler.py b/lmdeploy/pytorch/paging/prefill_scheduler.py index d2988974d6..79162faafe 100644 --- a/lmdeploy/pytorch/paging/prefill_scheduler.py +++ b/lmdeploy/pytorch/paging/prefill_scheduler.py @@ -16,7 +16,11 @@ from lmdeploy.messages import EventType from lmdeploy.pytorch import envs as _envs -from lmdeploy.pytorch.long_context import get_long_context_chunk_limit, plan_long_context_chunk +from lmdeploy.pytorch.long_context import ( + get_last_prefill_save_step, + get_long_context_chunk_limit, + plan_long_context_chunk, +) from lmdeploy.utils import get_logger from ..config import CacheConfig, SchedulerConfig @@ -145,7 +149,7 @@ def _get_reorder_info(self, seq: SchedulerSequence): return info def _estimate_long_chunks(self, seq: SchedulerSequence, chunk_limit: int, first_chunk_end: int): - """Count the planned first chunk and estimate its continuation in O(1). + """Estimate chunk count without walking the remaining text chunks. The continuation accounts for save alignment but treats future tokens as text. Later multimodal spans can change its boundaries, so this remains an ordering estimate, never an allocation or execution plan. @@ -154,6 +158,13 @@ def _estimate_long_chunks(self, seq: SchedulerSequence, chunk_limit: int, first_ end = seq.num_history_ids + seq.num_token_ids remaining = end - first_chunk_end ordinary_chunks = 1 + (remaining + chunk_limit - 1) // chunk_limit + if seq._seq_meta.prefill_save_last_only: + save_step = get_last_prefill_save_step(seq) + if save_step is None or save_step <= first_chunk_end: + return ordinary_chunks + before_save = (save_step - first_chunk_end + chunk_limit - 1) // chunk_limit + after_save = (end - save_step + chunk_limit - 1) // chunk_limit + return 1 + before_save + after_save alignment = seq._seq_meta.prefill_save_alignment if not alignment or len(seq.history_embeddings) > 0: return ordinary_chunks diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py index c6ea95f3ba..0d6aaf22c1 100644 --- a/lmdeploy/pytorch/paging/scheduler.py +++ b/lmdeploy/pytorch/paging/scheduler.py @@ -86,6 +86,7 @@ def __init__( seq_meta.prefill_save_alignment = ( cache_config.prefill_save_alignment if self.is_ssm and (cache_config.enable_prefix_caching or cache_config.num_store_state_caches) else 0) + seq_meta.prefill_save_last_only = bool(seq_meta.prefill_save_alignment) and cache_config.prefill_save_last_only seq_meta.prefix_cache_checkpoint_block_size = ( cache_config.block_size if cache_config.enable_prefix_caching and cache_config.prefix_cache_block_aligned else 0) diff --git a/tests/pytorch/engine/test_executor_base.py b/tests/pytorch/engine/test_executor_base.py index 99aa0fab14..4381f7338a 100644 --- a/tests/pytorch/engine/test_executor_base.py +++ b/tests/pytorch/engine/test_executor_base.py @@ -71,10 +71,15 @@ def test_mooncake_state_budget_is_deployable_and_excluded_from_runtime(slots): from lmdeploy.cli.utils import ArgumentHelper parser = ArgumentParser() ArgumentHelper.prefill_save_alignment(parser) + ArgumentHelper.prefill_save_last_only(parser) ArgumentHelper.mooncake_state_save_slots(parser) defaults = parser.parse_args([]) - assert (defaults.prefill_save_alignment, defaults.mooncake_state_save_slots) == (8192, 8) - args = parser.parse_args(['--mooncake-state-save-slots', str(slots), '--prefill-save-alignment', '4096']) + assert defaults.prefill_save_alignment == PytorchEngineConfig().prefill_save_alignment == 1024 + assert CacheConfig.prefill_save_alignment == 1024 + assert defaults.mooncake_state_save_slots == 8 + assert not defaults.prefill_save_last_only + args = parser.parse_args(['--mooncake-state-save-slots', str(slots), '--prefill-save-alignment', '4096', + '--prefill-save-last-only']) config = ConfigBuilder.build_cache_config(PytorchEngineConfig( max_batch_size=2, prefix_cache_state_budget=1, kv_transfer_config=KVTransferConfig(kv_connector='MooncakeStoreConnector', kv_role='kv_producer'), @@ -84,6 +89,7 @@ def test_mooncake_state_budget_is_deployable_and_excluded_from_runtime(slots): executor.cache_config = config size = executor._get_state_cache_mem() assert config.prefill_save_alignment == 4096 + assert config.prefill_save_last_only assert config.num_state_caches == 5 + slots assert size == config.num_state_caches * StateCacheEngine.get_state_slot_nbytes(config.states_shapes) manager = build_state_manager(config) diff --git a/tests/pytorch/kv_connector/test_mooncake_store_scheduler.py b/tests/pytorch/kv_connector/test_mooncake_store_scheduler.py index dd5de14183..8984cfb3e6 100644 --- a/tests/pytorch/kv_connector/test_mooncake_store_scheduler.py +++ b/tests/pytorch/kv_connector/test_mooncake_store_scheduler.py @@ -28,6 +28,7 @@ ) from lmdeploy.pytorch.kv_connector.mooncake.store.scheduler import MooncakeStoreScheduler from lmdeploy.pytorch.kv_connector.mooncake.store.worker import MooncakeStoreWorker +from lmdeploy.pytorch.long_context import get_last_prefill_save_step, plan_long_context_chunk from lmdeploy.pytorch.messages import MessageStatus, SequenceMeta from lmdeploy.pytorch.model_inputs import ModelInputs from lmdeploy.pytorch.multimodal.data_type import MultiModalData @@ -140,11 +141,17 @@ def step(end): scheduler.shutdown() -@pytest.mark.parametrize(('lookup_alignment', 'expected_hit'), [(8, 24), (12, 24), (16, 16)]) -def test_hybrid_multimodal_save_lookup_and_load_without_local_prefix_cache(lookup_alignment, expected_hit): - cache_config = _hybrid_cache_config() +@pytest.mark.parametrize(('last_only', 'lookup_alignment', 'mtp', 'expected_hit'), [ + (False, 8, False, 24), (False, 12, False, 24), (False, 16, False, 16), + (True, 8, False, 24), (True, 16, False, 0), + (False, 8, True, 16), (True, 8, True, 0), +]) +def test_hybrid_multimodal_save_lookup_and_load_without_local_prefix_cache( + last_only, lookup_alignment, mtp, expected_hit): + cache_config = replace(_hybrid_cache_config(), prefill_save_last_only=last_only) + strategy = ARSpecSequenceStrategy() if mtp else ARSequenceStrategy() paging = Scheduler(SchedulerConfig(max_batches=1, max_session_len=64), cache_config, - seq_meta=SequenceMeta(4, strategy=ARSequenceStrategy())) + seq_meta=SequenceMeta(4, strategy=strategy)) producer = MooncakeStoreScheduler(cache_config) image = MultiModalData(torch.tensor([1]), start=2, end=10) source = paging.add_session(0).add_sequence(range(25), multimodals={'image': [image]}) @@ -156,7 +163,7 @@ def test_hybrid_multimodal_save_lookup_and_load_without_local_prefix_cache(looku (source,), (end,), (blocks,), (blocks,), (1,))) # The first boundary is inside the image; later saves cover its FA # blocks but only persist state at the exact safe boundaries. - if end == 8: + if end == 8 or (last_only and end != 24): assert result is None continue save = result.save_requests[0] @@ -177,12 +184,15 @@ def test_hybrid_multimodal_save_lookup_and_load_without_local_prefix_cache(looku side_effect=lambda req_id, token_len, hashes, non_block, recompute_blocks: worker.lookup( token_len, hashes, recompute_blocks=recompute_blocks)) consumer_paging = Scheduler(SchedulerConfig(max_batches=1, max_session_len=64), config, - seq_meta=SequenceMeta(4, strategy=ARSequenceStrategy()), kv_connector=consumer) + seq_meta=SequenceMeta(4, strategy=strategy), kv_connector=consumer) assert not consumer_paging.block_trie.enabled for session_id, (content, span_end, expected) in enumerate( ((1, 10, expected_hit), (2, 10, 0), (1, 11, 0)), start=1): image = MultiModalData(torch.tensor([content]), start=2, end=span_end) - request = consumer_paging.add_session(session_id).add_sequence(range(25), multimodals={'image': [image]}) + # With block rewind, even an appended prompt cannot reuse state at 24 + # without FA through 28; last-only saving has no earlier state fallback. + token_ids = range(29 if mtp else 25) + request = consumer_paging.add_session(session_id).add_sequence(token_ids, multimodals={'image': [image]}) assert request.prefix_cache.multimodal_spans assert consumer.get_num_new_matched_tokens(request, 0) == (expected, expected > 0) if expected: @@ -204,6 +214,75 @@ def test_hybrid_multimodal_save_lookup_and_load_without_local_prefix_cache(looku paging.shutdown() +@pytest.mark.parametrize('mtp', [False, True]) +def test_last_only_prefill_saves_both_caches_and_reuses_across_turns(mtp): + cache = replace(_hybrid_cache_config(), num_gpu_blocks=24, num_state_caches=12, + enable_prefix_caching=True, prefix_cache_block_aligned=True, + prefix_cache_token_lookahead=int(mtp), prefill_save_last_only=True, + max_prefill_token_num=8) + strategy = ARSpecSequenceStrategy(recompute_blocks=0) if mtp else ARSequenceStrategy() + paging = Scheduler(SchedulerConfig(max_batches=1, max_session_len=64), cache, + seq_meta=SequenceMeta(4, strategy=strategy)) + store = MooncakeStoreScheduler(cache) + checkpoints = paging.block_trie.state_checkpoints + multimodals = {'image': [MultiModalData(torch.tensor([1]), start=24, end=25)]} + seq = paging.add_session(0).add_sequence(range(25), multimodals=multimodals) + paging.state_manager.allocate(seq) + published = [] + for expected_step in (16 if mtp else 24, 40): + paging.block_manager.allocate(seq) + paging.block_trie.allocate(seq) + assert get_last_prefill_save_step(seq) == expected_step + saved = [] + while seq.num_token_ids: + end = plan_long_context_chunk(seq, cache.max_prefill_token_num).chunk_end + local = checkpoints.reserve_prefill_save_batch([seq], (end,)) + blocks = tuple(range(end // cache.block_size)) + meta = store.build_connector_meta(_connector_step( + (seq,), (end,), (blocks,), (blocks,), (seq.logical_state,))) + saves = () if meta is None else meta.save_requests + if end == expected_step: + save, = saves + assert save.start_block == 0 + assert save.block_ids == blocks + assert len(save.block_hashes) == len(blocks) + assert save.state.boundary_tokens == end + assert local.state_pairs == ((seq.logical_state, seq.prefix_cache.pending_save.slot),) + assert save.state.source_slot == seq.logical_state + assert save.state.snapshot_slot != seq.prefix_cache.pending_save.slot + published.append(seq.prefix_cache.pending_save.node) + assert checkpoints.publish_save(seq) + saved.append(save.save_id) + else: + assert not local.state_pairs + assert not saves + seq.set_step(end) + assert len(saved) == 1 + assert all(node.state_checkpoint.published for node in published) + store.update_connector_output(KVConnectorOutput(completed_save_ids=set(saved))) + assert store._state_slots.get_num_free() == cache.mooncake_state_save_slots + if seq.input_end_pos == 25: + seq.update_token_ids(range(25, 41)) + store.on_new_request(seq) + + # A hit at this turn's only save boundary still restores state, without + # reserving another checkpoint for its remaining input token. + paging.end_session(0) + warm = paging.add_session(1).add_sequence(range(41), multimodals=multimodals) + paging.block_trie.match(warm) + assert warm.num_history_ids == 40 + paging.state_manager.allocate(warm) + restore = checkpoints.prepare_restore_batch([warm]) + assert restore.state_pairs == ((published[-1].state_checkpoint.slot, warm.logical_state),) + assert not checkpoints.reserve_prefill_save_batch([warm], (41,)).state_pairs + blocks = tuple(range(10)) + assert store.build_connector_meta(_connector_step( + (warm,), (41,), (blocks,), (blocks,), (warm.logical_state,))) is None + checkpoints.finish_forward_dispatch([warm], has_save_plan=False) + paging.shutdown() + store.shutdown() + + @pytest.mark.parametrize('invalid_slot', [-1, 0, 3, 4]) def test_hybrid_load_binding_rejects_non_runtime_slots_without_consuming_lookup(invalid_slot): connector = MooncakeStoreScheduler(_hybrid_cache_config('kv_both')) diff --git a/tests/pytorch/paging/test_prefill_scheduler.py b/tests/pytorch/paging/test_prefill_scheduler.py index c5bc1f177c..305848797c 100644 --- a/tests/pytorch/paging/test_prefill_scheduler.py +++ b/tests/pytorch/paging/test_prefill_scheduler.py @@ -9,7 +9,7 @@ import lmdeploy.pytorch.paging.prefill_scheduler as prefill_scheduler_module from lmdeploy.messages import KVTransferConfig from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig -from lmdeploy.pytorch.long_context import plan_long_context_chunk +from lmdeploy.pytorch.long_context import get_last_prefill_save_step, plan_long_context_chunk from lmdeploy.pytorch.messages import InputEmbeddings, MessageStatus, SequenceMeta, UpdateTokenMode from lmdeploy.pytorch.multimodal.data_type import MultiModalData from lmdeploy.pytorch.paging.scheduler import Scheduler @@ -476,19 +476,21 @@ def test_multimodal_chunk_reservation_obeys_the_draft_input_shift(method): (False, True, True, [800]), (True, True, True, [8192, 15872, 16384]), ]) -def test_prefill_state_save_alignment_is_shared(hybrid, prefix_caching, store, ends): +@pytest.mark.parametrize('last_only', [False, True]) +def test_prefill_state_save_alignment_is_shared(hybrid, prefix_caching, store, ends, last_only): from lmdeploy.pytorch.engine.inputs_maker import LongContextChunker from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy cache = CacheConfig( max_batches=2, block_size=64, num_cpu_blocks=0, num_gpu_blocks=256, - max_prefill_token_num=8192, prefill_save_alignment=512, + max_prefill_token_num=8192, prefill_save_alignment=512, prefill_save_last_only=last_only, enable_prefix_caching=prefix_caching, prefix_cache_block_aligned=hybrid, states_shapes=[((2,), torch.float32)] if hybrid else [], num_state_caches=16, kv_transfer_config=(KVTransferConfig(kv_connector='MooncakeStoreConnector', kv_role='kv_both') if store else None)) scheduler = Scheduler(SchedulerConfig(max_batches=2, max_session_len=32768), cache, seq_meta=SequenceMeta(64, strategy=ARSequenceStrategy())) + assert scheduler.seq_meta.prefill_save_last_only == (last_only and hybrid and (prefix_caching or store)) seq = scheduler.add_session(0).add_sequence(range(ends[-1])) prefill = scheduler._prefill_scheduler info = prefill_scheduler_module._PrefillReorderer(prefill)._get_reorder_info(seq) @@ -505,6 +507,49 @@ def test_prefill_state_save_alignment_is_shared(hybrid, prefix_caching, store, e scheduler.shutdown() +@pytest.mark.parametrize(('history', 'budget', 'spans', 'mtp', 'save_step', 'ends'), [ + (0, 6, (), False, 16, (6, 12, 16, 20)), + (0, 32, (), False, 16, (16, 24)), + (0, 32, (), False, None, (8,)), + (0, 32, (), False, None, (7,)), + (16, 6, (), False, 16, (20,)), + (17, 6, (), False, 16, (20,)), + (0, 32, ((14, 18),), False, 8, (8, 20)), + (0, 32, ((6, 10), (14, 18)), False, None, (20,)), + (0, 6, ((14, 18),), False, 8, (6, 8, 14, 20)), + (0, 32, ((16, 18),), True, 8, (8, 20)), +]) +def test_last_only_prefill_boundaries(history, budget, spans, mtp, save_step, ends): + from lmdeploy.pytorch.engine.inputs_maker import LongContextChunker + from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy + from lmdeploy.pytorch.strategies.ar_spec.sequence import ARSpecSequenceStrategy + + strategy = ARSpecSequenceStrategy(recompute_blocks=0) if mtp else ARSequenceStrategy() + cache = CacheConfig( + max_batches=1, block_size=4, num_cpu_blocks=0, num_gpu_blocks=16, + max_prefill_token_num=budget, states_shapes=[((2,), torch.float32)], num_state_caches=6, + prefill_save_alignment=8, prefill_save_last_only=True, mooncake_state_save_slots=2, + kv_transfer_config=KVTransferConfig(kv_connector='MooncakeStoreConnector', kv_role='kv_producer')) + scheduler = Scheduler(SchedulerConfig(max_batches=1, max_session_len=64), cache, + seq_meta=SequenceMeta(4, strategy=strategy)) + multimodals = {'image': [MultiModalData(torch.tensor([1]), *span) for span in spans]} + seq = scheduler.add_session(0).add_sequence(range(ends[-1]), multimodals=multimodals) + seq.set_step(history) + prefill = scheduler._prefill_scheduler + info = prefill_scheduler_module._PrefillReorderer(prefill)._get_reorder_info(seq) + assert info.estimated_long_chunks == len(ends) + assert info.prefill_token_count == ends[0] - history + chunker = LongContextChunker(budget) + chunker.set_seq(seq) + for end in ends: + assert get_last_prefill_save_step(seq) == save_step + assert prefill._next_long_context_chunk_end(seq) == end + assert chunker.next_chunk_size()[0] == end - seq.num_history_ids + assert chunker.is_last_chunk() == (end == ends[-1]) + seq.set_step(end) + scheduler.shutdown() + + def _make_mooncake_prefill_scheduler(budget: int = 32, hybrid: bool = True): from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy