Ssm prefix cache non aligned

[grimoire]

Requirement

• optimize and modularize SSM prefix caching #4788
• [Bugfix] Fix PyTorch H2D input lifetime across CUDA streams #4792

Motivation

SSM prefix caching currently saves and restores recurrent-state checkpoints only at scheduler block boundaries. This prevents reuse when prefill or chunked prefill ends at a safe but non-block-aligned step, which is common for long-context and multimodal requests.

This PR allows SSM checkpoints at exact prefill boundaries while preserving the existing full-block trie ownership model and keeping decode checkpoints block-aligned.

This PR is temporarily based on ssm-prefix-cache-readability and will be rebased onto main after that branch is merged.

Modification

• Store the exact checkpoint step and an optional frozen partial KV block on each checkpoint record.
• Keep trie nodes responsible only for canonical full blocks. A non-aligned checkpoint copies its partial KV block into checkpoint-owned storage.
• Restore the frozen partial block into a request-owned writable block before model execution.
• Add ordered logical KV-block copy support for packed cache pools, including scheduler-block to kernel-block geometry.
• Carry KV and recurrent-state restore/save plans through a dedicated one-forward CacheCheckpointInputs object.
• Execute restore copies before the model forward and save copies afterward on the forward stream.
• Index candidates by adapter, exact step, and partial-tail hash, followed by full token, multimodal identity, block-path, and ownership verification.
• Pin checkpoints while asynchronous restore or save copies are pending.
• Release frozen partial blocks through checkpoint lifecycle management and prefer them during KV-pressure eviction.
• Keep partial hits disabled when routed-expert tail history cannot be reconstructed.
• Preserve writable suffix blocks for normal continuation and MTP recompute overlap.

No public API or configuration behavior is changed.

BC-breaking

No.

Use cases

• Reuse an SSM prefix ending at steps such as B - 1, B + 1, or other safe prefill boundaries.
• Cache non-block-aligned long-context chunks.
• Cache multimodal prefill checkpoints when the boundary is outside multimodal spans.
• Restore the same frozen checkpoint concurrently into multiple request-owned blocks.

[Copilot]

Pull request overview

This PR extends the PyTorch SSM prefix-cache implementation to support non-block-aligned prefill checkpoint boundaries while keeping trie KV ownership block-aligned. It introduces a per-forward CacheCheckpointInputs object that carries ordered KV/state restore/save copy plans, adds logical KV block-copy support (including packed cache pools), and updates the SSM checkpoint index/lifecycle to track exact steps plus optional frozen partial KV tails.

Changes:

• Add support for SSM checkpoints at exact (non-block-aligned) prefill boundaries via frozen partial KV tail blocks owned by checkpoints.
• Introduce CacheCheckpointInputs and route KV/state restore+save plans through InputsMaker → ModelAgent, executing copies around the model forward on the forward stream.
• Add logical KV-block copy implementations (default + CUDA/Triton) and update cache engine/backends/tests accordingly.

Reviewed changes

Copilot reviewed 37 out of 37 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/pytorch/paging/test_scheduler.py	Updates scheduler test to use new state save-plan helper.
tests/pytorch/paging/test_block_trie/test_trie.py	Updates trie cursor attachment assertion to new helper.
tests/pytorch/paging/test_block_trie/test_node.py	Updates node checkpoint construction and attachment/root helpers.
tests/pytorch/paging/test_block_trie/test_checkpoint.py	Adds partial-checkpoint matching tests and new checkpoint geometry helpers coverage.
tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py	Expands lifecycle coverage for partial tails, replacement, and eviction behavior.
tests/pytorch/paging/test_block_manager.py	Adds coverage for resolving logical allocator IDs to GPU block offsets.
tests/pytorch/kernel/test_copy_packed_cache.py	New CUDA/Triton test coverage for packed-cache logical block copies and CacheEngine dispatch.
tests/pytorch/engine/test_model_agent.py	Adds ordering and lifetime tests for checkpoint KV/state copies through ModelAgent.
tests/pytorch/engine/test_inputs_maker.py	Adds tests for KV copy plan creation and new prefill/decode cache-input preparation paths.
tests/pytorch/engine/test_cache_inputs.py	New tests for CacheCheckpointInputs device transfer + stream recording semantics.
tests/pytorch/engine/test_cache_engine.py	Updates state-cache copy API expectations (tuple-based indices).
tests/pytorch/engine/test_cache_block_copy.py	New tests for default backend KV logical-block copy and CacheEngine validation/dispatch.
lmdeploy/pytorch/strategies/dllm/step_inputs.py	Removes deprecated decode checkpoint fields from DLLM reindex path.
lmdeploy/pytorch/strategies/ar/step_inputs.py	Removes deprecated decode checkpoint fields from AR reindex path.
lmdeploy/pytorch/strategies/ar/model_inputs.py	Drops unused index-select parameters tied to removed checkpoint fields.
lmdeploy/pytorch/strategies/ar_spec/step_inputs.py	Removes deprecated decode checkpoint fields from AR-spec reindex path.
lmdeploy/pytorch/paging/scheduler.py	Adds scheduler helper to resolve logical block IDs to GPU offsets for copy planning.
lmdeploy/pytorch/paging/block_trie/trie.py	Implements partial checkpoint matching/apply semantics and integrates frozen-tail eviction.
lmdeploy/pytorch/paging/block_trie/README.md	Updates design/ownership documentation to reflect partial checkpoints + CacheCheckpointInputs.
lmdeploy/pytorch/paging/block_trie/node.py	Extends node checkpoint metadata with exact step and optional frozen_block_id.
lmdeploy/pytorch/paging/block_trie/checkpoint.py	Changes sparse key to (adapter, step, tail_hash) and adds partial-step geometry helpers.
lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py	Adds frozen-tail reservation/release and eviction path; supports exact-step reservations.
lmdeploy/pytorch/paging/block_manager/base_block_manager.py	Adds resolve_gpu_block_offsets() with validation for copy-plan construction.
lmdeploy/pytorch/model_inputs.py	Removes old SSM checkpoint copy fields from ModelInputs/ModelInputsDelta.
lmdeploy/pytorch/messages.py	Adds is_prefix_cache_boundary_safe() for exact-step boundary validation.
lmdeploy/pytorch/kernels/cuda/copy_packed_cache.py	New Triton kernel to copy packed logical KV blocks with one launch.
lmdeploy/pytorch/engine/model_agent/agent.py	Routes restore/save via CacheCheckpointInputs and executes KV/state copies around forward.
lmdeploy/pytorch/engine/inputs_maker.py	Builds per-forward CacheCheckpointInputs (KV + state plans) for prefill/chunk/decode.
lmdeploy/pytorch/engine/engine_loop.py	Switches save-detection logic to cache_inputs.state_save_plan.
lmdeploy/pytorch/engine/cache_inputs.py	New CacheCheckpointInputs dataclass with device-transfer and stream-recording support.
lmdeploy/pytorch/engine/cache_engine.py	Adds logical KV-block copy build + copy_logical_blocks() validation/dispatch.
lmdeploy/pytorch/backends/default/op_backend.py	Registers default OpType.CacheBlockCopy builder.
lmdeploy/pytorch/backends/default/cache_block_copy.py	New default (torch) logical-block copy implementation with chunked workspaces.
lmdeploy/pytorch/backends/cuda/op_backend.py	Registers CUDA OpType.CacheBlockCopy builder.
lmdeploy/pytorch/backends/cuda/cache_block_copy.py	New CUDA packed-cache copy implementation using Triton kernel launcher.
lmdeploy/pytorch/backends/cache_block_copy.py	New backend interface for cache-block-copy builder/impl.
lmdeploy/pytorch/backends/base.py	Adds new backend op enum CacheBlockCopy.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[yidingcheng0206]

Non-blocking question: the new frozen partial-KV restore relies on the existing restore-pin lifetime. Is it guaranteed that the restore copy has already been queued on the forward stream before unpin_restores() allows prefetch eviction to reuse the frozen block?

[grimoire]

Non-blocking question: the new frozen partial-KV restore relies on the existing restore-pin lifetime. Is it guaranteed that the restore copy has already been queued on the forward stream before unpin_restores() allows prefetch eviction to reuse the frozen block?

Any requests to the ModelAgent is FIFO ,while ModelAgent processes forwards serially and queues the restore copy before the current model execution. Thus the restore read is ordered before a later overwrite of the recycled block.