optimize and modularize SSM prefix caching

[grimoire]

Motivation

SSM prefix-cache hits require both reusable KV blocks and an exact recurrent-state checkpoint.

The existing aligned checkpoint matcher first found a sparse candidate, then reconstructed and verified its complete trie ancestry in Python for every request. At large batch sizes and long contexts, this CPU work could take hundreds of milliseconds or seconds, making prefix caching unable to overlap effectively with GPU execution.

This PR accelerates the existing aligned SSM matcher and reorganizes the prefix-cache implementation around explicit ownership boundaries. It intentionally does not add non-block-aligned checkpoints.

Changes

Faster exact SSM checkpoint matching

• Store immutable contiguous token IDs, logical block IDs, and multimodal identity when a checkpoint is published.
• Keep sparse lookup partitioned by adapter, checkpoint step, and block identity.
• Perform exact verification with bounded Python work and NumPy comparisons instead of walking every trie block in Python.
• Preserve collision handling, stale-entry cleanup, multimodal identity checks, MTP recompute overlap, routed-expert replay, and tentative-match rollback.
• Keep the implementation Python-only and introduce no device synchronization.

Clearer prefix-cache ownership

Convert block_trie.py into a package with components organized by owned state:

• trie.py: prefix identity, matching, allocation, and routed-expert replay.
• checkpoint.py: sparse checkpoint lookup and exact verification.
• checkpoint_lifecycle.py: state-slot reservation, publication, pins, eviction, and stale cleanup.
• kv_lifecycle.py: trie KV references and leaf eviction.
• node.py: monotonic trie topology and lazy node checkpoint ownership.
• prefix_cache_state.py: per-sequence restore, save, producer-pin, multimodal, and recompute-overlap state.

Additional readability changes include:

• Make tentative scheduler admission and rollback ownership explicit.
• Narrow BlockTrie construction to its actual dependencies.
• Separate block indices from allocator block IDs in naming.
• Use consistent checkpoint reserve/publish/pin/release vocabulary.
• Remove redundant wrappers and repeated prefix-cache enable checks.
• Split tests by component ownership.
• Add paging/block_trie/README.md describing the pipeline, ownership map, lifecycle, and invariants.

Performance

Production-shaped aligned exact-hit benchmarks used one pinned CPU, fixed Python hashing, one numerical worker thread, two warmups, and ten measured samples.

The optimized BlockTrie.match() scope includes lookup, exact verification, allocator access/refcount updates, logical-block attachment, sequence-step mutation, and match statistics.

Batch	Context	Previous p99	Optimized p50	Optimized p99
128	64K	~675 ms	8.85 ms	9.41 ms
512	64K	~2.67 s	48.13 ms	68.19 ms
128	128K	~1.37 s	17.73 ms	18.60 ms
512	128K	~5.23 s	95.72 ms	120.81 ms

Repeated measurements throughout the readability refactor showed no meaningful SSM matcher regression. Controlled AR matching and allocation measurements also showed no regression.

Memory trade-off

Each published SSM checkpoint retains immutable host-side exact-match metadata for its prefix. In the worst tested configuration—512 distinct 128K checkpoints—the additional token and logical-block identity data is approximately 520 MiB.

The metadata lifetime is bounded by the existing checkpoint budget and eviction lifecycle. This trade-off removes the dominant per-hit Python trie reconstruction without adding native code or a second checkpoint store.

[Copilot]

Pull request overview

This PR reorganizes PyTorch prefix caching (KV trie + SSM state checkpoints) into smaller ownership-focused modules and accelerates aligned SSM checkpoint matching by publishing immutable exact-match metadata for fast verification, reducing Python overhead on prefix-cache hits.

Changes:

• Replaces the monolithic block_trie.py with a block_trie/ package (trie, node topology, KV lifecycle, checkpoint lookup + lifecycle, docs).
• Introduces explicit per-sequence protocol state in prefix_cache_state.py and wires it through scheduler/input maker/engine loop.
• Updates and expands tests to reflect the new APIs and lifecycle semantics (reserve/publish/pin/unpin, recompute overlap behavior).

Reviewed changes

Copilot reviewed 26 out of 27 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 tests to new prefix-cache and checkpoint lifecycle APIs (state_checkpoints, trie_cursor, restore/pending_save).
tests/pytorch/paging/test_block_trie/test_trie.py	Adds detailed trie/match/allocate/recompute-overlap tests for the modularized BlockTrie.
tests/pytorch/paging/test_block_trie/test_node.py	Adds unit tests for monotonic node topology and lazy checkpoint record allocation.
tests/pytorch/paging/test_block_trie/test_kv_lifecycle.py	Adds eviction and stale-leaf handling tests for KV lifecycle behavior.
tests/pytorch/paging/test_block_trie/test_checkpoint.py	Adds tests for sparse SSM checkpoint lookup and exact verification semantics.
tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py	Adds extensive tests for reserve/publish/pin/unpin/evict and publication rollback behavior.
tests/pytorch/paging/test_block_trie/conftest.py	Adds fixtures for block trie and SSM scheduler configurations used by the new component tests.
tests/pytorch/paging/test_block_trie/_utils.py	Adds shared test helpers and an SSM checkpoint publication helper.
tests/pytorch/paging/test_block_trie/init.py	Initializes the new test_block_trie test package.
tests/pytorch/engine/test_inputs_maker.py	Updates engine-loop checkpoint pin/unpin tests to use state_checkpoints APIs and new restore representation.
lmdeploy/pytorch/strategies/ar_spec/sequence.py	Switches AR-spec strategy to configure recompute overlap via prefix_cache.recompute_overlap.
lmdeploy/pytorch/prefix_cache_state.py	New module defining per-sequence prefix-cache protocol state (restore/save/pins, recompute overlap, multimodal identity).
lmdeploy/pytorch/paging/seq_states/states.py	Updates sequence-free path for new prefix-cache fields and checkpoint lifecycle calls.
lmdeploy/pytorch/paging/scheduler.py	Refactors prefix-cache admission/rollback flow; updates BlockTrie construction and stats handling; new gate-match structure.
lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py	Updates eviction logic to use block_trie.enabled and state_checkpoints.evict().
lmdeploy/pytorch/paging/block_trie/trie.py	New BlockTrie facade implementing KV matching/allocation, SSM checkpoint matching, and recompute-overlap policy.
lmdeploy/pytorch/paging/block_trie/README.md	New documentation describing pipeline, ownership map, invariants, and where to make changes.
lmdeploy/pytorch/paging/block_trie/node.py	New monotonic trie node implementation with lazy NodeStateCheckpoint record.
lmdeploy/pytorch/paging/block_trie/kv_lifecycle.py	New KV reference transaction + leaf eviction component.
lmdeploy/pytorch/paging/block_trie/checkpoint.py	New sparse checkpoint index + exact verification primitives and match metadata freezing.
lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py	New checkpoint lifecycle manager (reserve/publish/pin/unpin/evict + stale cleanup).
lmdeploy/pytorch/paging/block_trie/init.py	Exposes the new BlockTrie package public API.
lmdeploy/pytorch/paging/block_trie.py	Removes the old monolithic BlockTrie implementation.
lmdeploy/pytorch/messages.py	Moves prefix-cache state types to prefix_cache_state.py, updates multimodal identity indexing, and renames identity accessor.
lmdeploy/pytorch/engine/inputs_maker.py	Updates restore/save planning to use the new state_checkpoints lifecycle and restore state object.
lmdeploy/pytorch/engine/engine_loop.py	Updates checkpoint publication/unpin boundaries to use new state_checkpoints APIs.

---

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