refactor(pytorch): redesign CacheEngine around plans and allocations

[grimoire]

Motivation

CacheEngine had accumulated several conflicting responsibilities:

• deriving cache tensor shapes and dtypes;
• selecting physical packing and layout;
• allocating storage;
• constructing model-facing cache views;
• moving caches through swap, local copy, and PD migration;
• supporting backend-specific behavior through dlinfer monkey patches.

New cache users—including DSA index caches, DeepSeek V4 compressed caches, and heterogeneous backend layouts—made the old configuration-owned, single-pool design increasingly difficult to extend safely.

This PR redesigns CacheEngine around one worker-local pipeline:

built cache-using operators
        |
        v
BlockCacheRequest[]
        |
        v
BlockCachePlan
        |
        v
backend-selected layout
        |
        v
CacheAllocation
        |
        v
CacheEngine runtime lifecycle

The plan used for memory sizing is also used for real CPU and accelerator allocation, so layout selection and memory accounting have one source of truth.

What changed

Read README.md in pytorch/engine/cache_engine for more details.

[Copilot]

Pull request overview

This PR refactors the PyTorch cache engine into a plan/layout/engine architecture, shifting model-specific “extra” cache allocation from static config (block_cache_specs) to operator-declared block-cache requests collected from built modules. It also updates executors/agents to build and retain worker-local cache plans before cache-engine construction, and adjusts DeepSeek V4 + DSA indexer/compressor codepaths to use the new operator-owned block-cache views.

Changes:

• Introduce lmdeploy.pytorch.engine.cache_engine package (schema/plan/layout/view/state/engine/migration/collector) and wire it through executors and model agents.
• Replace legacy packed-cache copy kernels/builders with allocation-aware block copy via backend CacheBackend (CUDA Triton path + torch fallback).
• Update DeepSeek V4 + DSA indexer/compressor/indexer interfaces to request and resolve block caches via the new binding/NamedCacheView mechanism; update and reorganize tests accordingly.

Reviewed changes

Copilot reviewed 76 out of 76 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/test_lmdeploy/test_fp8_kv_cache_policy.py	Updates tests to use new cache-engine schema helpers for quant cache descs.
tests/pytorch/paging/test_block_manager.py	Consolidates GPU-offset validation tests into a single “unavailable blocks” test.
tests/pytorch/kernel/test_v4_compressor.py	Adapts V4 compressor tests to new block_caches mapping and is_indexer flag.
tests/pytorch/kernel/test_copy_packed_cache.py	Removes legacy packed-cache copy tests tied to the old copy primitive.
tests/pytorch/kernel/test_copy_cache.py	Adds tests for new contiguous-pool Triton cache-block copy kernel.
tests/pytorch/engine/test_model_agent.py	Adds coverage for worker-local cache plan retention, graph-runner request collection, checkpoint copy API changes, and level-2 wakeup ordering/validation.
tests/pytorch/engine/test_executor_base.py	Updates executor tests for new cache-plan sizing flow and sparse-MLA cache policy finalization.
tests/pytorch/engine/test_cache_engine/test_view.py	Adds unit tests for NamedCacheView mapping/row/layer resolution behavior.
tests/pytorch/engine/test_cache_engine/test_state_engine.py	Adds unit tests for new StateCacheEngine allocation + slot copy/zero semantics.
tests/pytorch/engine/test_cache_engine/test_schema.py	Adds unit tests for schema types (CacheDesc/geometry/requests/specs).
tests/pytorch/engine/test_cache_engine/test_plan.py	Adds unit tests for BlockCachePlan validation, allocation, and metadata.
tests/pytorch/engine/test_cache_engine/test_layout.py	Adds unit tests for layouts (packed/row/contiguous/composite) and default backend selection.
tests/pytorch/engine/test_cache_engine/test_copy.py	Adds unit tests for backend block-copy primitives + CacheEngine copy dispatch/validation.
tests/pytorch/engine/test_cache_engine/test_collector.py	Adds unit tests for module request collection and consumer-row binding.
tests/pytorch/engine/test_cache_engine/init.py	Creates package init for the new focused cache-engine test suite.
tests/pytorch/engine/test_cache_engine.py	Removes large legacy cache-engine test file in favor of split suite.
tests/pytorch/engine/test_cache_block_copy.py	Removes tests for removed CacheBlockCopy op-based copy path.
tests/pytorch/config/test_model_config.py	Removes DeepSeek V4 block-cache-spec normalization checks and post-build hook usage.
lmdeploy/pytorch/spec_decode/spec_agent.py	Adds draft cache-plan construction/retention and passes finalized plans into CacheEngine.
lmdeploy/pytorch/spec_decode/proposers/base.py	Removes outdated comment tied to config-declared cache specs.
lmdeploy/pytorch/spec_decode/base.py	Adds base build_cache_plan hook and stores block_cache_plan.
lmdeploy/pytorch/paging/block_trie/README.md	Updates SSM checkpoint docs to reflect partial frozen-block handling and new copy-contract wording.
lmdeploy/pytorch/nn/v4_indexer.py	Switches indexer interface to accept a named block_caches mapping.
lmdeploy/pytorch/nn/v4_compressor.py	Adds block-cache request/binding lifecycle and resolves bound caches per forward; updates write interface.
lmdeploy/pytorch/nn/v4_attention.py	Updates attention forward signature to accept window_state_fp8 + block_caches.
lmdeploy/pytorch/nn/nsa.py	Moves DSA indexer cache acquisition to bound operator-owned block caches via StepContext.
lmdeploy/pytorch/models/glm_moe_dsa.py	Updates model to new DSA indexer API (no explicit cache tensor arguments).
lmdeploy/pytorch/models/deepseek_v4.py	Removes config-declared block caches; resolves operator-owned block caches via compressor bindings and passes mapping to attention/indexer.
lmdeploy/pytorch/models/deepseek_v32.py	Updates DSA indexer usage to new bound-cache path and removes cache-layer remapping.
lmdeploy/pytorch/model_inputs.py	Clarifies StepContext comments for block/state named cache views.
lmdeploy/pytorch/memdecode/agent.py	Adds memory-model cache-plan construction/retention and passes plan into CacheEngine.
lmdeploy/pytorch/kernels/cuda/copy_packed_cache.py	Removes old Triton packed-cache logical-block copy kernel.
lmdeploy/pytorch/kernels/cuda/copy_cache.py	Adds new Triton kernel to copy scheduler-sized blocks in contiguous cache pools across varying entry axes.
lmdeploy/pytorch/engine/model_agent/agent.py	Integrates plan building, passes model_config explicitly into model_forward, updates swap event waiting, and updates state checkpoint copy API (copy_slots).
lmdeploy/pytorch/engine/executor/uni_executor.py	Implements worker-local cache-plan preparation for single-process executor.
lmdeploy/pytorch/engine/executor/ray_executor.py	Implements distributed worker cache-plan preparation via collective RPC.
lmdeploy/pytorch/engine/executor/mp_executor.py	Implements multi-process worker cache-plan preparation via collective RPC.
lmdeploy/pytorch/engine/executor/base.py	Replaces old cache-block-size logic with per-worker plan sizes; validates rank list lengths; updates state-cache memory sizing call.
lmdeploy/pytorch/engine/executor/base_worker.py	Adds build_cache_plans RPC endpoint for workers.
lmdeploy/pytorch/engine/executor/init.py	Adds _finalize_sparse_mla_cache_policy and runs it before executor build/copy-to-worker.
lmdeploy/pytorch/engine/cache_inputs.py	Updates docstring to match new per-forward copy-plan semantics and host-vs-device plan ownership.
lmdeploy/pytorch/engine/cache_engine/view.py	Introduces NamedCacheView mapping for name/row/layer resolution across physical tensors.
lmdeploy/pytorch/engine/cache_engine/state.py	Introduces new StateCacheEngine implementation with allocation and host-indexed slot copy/zero semantics.
lmdeploy/pytorch/engine/cache_engine/schema.py	Introduces cache schema/types (CacheDesc, geometry, requests/bindings, tensor specs) and standard KV cache descriptor builders.
lmdeploy/pytorch/engine/cache_engine/README.md	Adds design documentation for the new cache-engine package and construction pipeline.
lmdeploy/pytorch/engine/cache_engine/plan.py	Introduces BlockCachePlan and plan builder integrating collected operator requests + backend layout selection.
lmdeploy/pytorch/engine/cache_engine/migration.py	Adds PD cache-pool metadata description and byte-transfer planning for multi-pool layouts.
lmdeploy/pytorch/engine/cache_engine/layout.py	Introduces physical allocation layouts and owning-pool metadata (CachePool, CacheAllocation, packed/row/contiguous/composite layouts).
lmdeploy/pytorch/engine/cache_engine/engine.py	Introduces new CacheEngine implementation owning allocations, swap pairs, local block copy, and PD migration wiring.
lmdeploy/pytorch/engine/cache_engine/collector.py	Introduces module request collection and consumer-row binding.
lmdeploy/pytorch/engine/cache_engine/init.py	Exposes new cache-engine public API surface (CacheEngine/CacheDesc/StateCacheEngine).
lmdeploy/pytorch/disagg/conn/protocol.py	Adds DistServeCachePoolInfo and includes cache pool metadata in PD endpoint info.
lmdeploy/pytorch/consts.py	Adds standardized V4 cache name constants and packed-token dimension constant.
lmdeploy/pytorch/configurations/deepseek_v4.py	Removes V4 block-cache post-build spec materialization; keeps state-cache specs and shared constants.
lmdeploy/pytorch/configurations/deepseek_v32.py	Removes config-based DSA index cache specs; relies on operator requests instead.
lmdeploy/pytorch/config.py	Removes BlockCacheSpec and related ModelConfig fields/hook; keeps named state-cache specs.
lmdeploy/pytorch/backends/nsa.py	Extends NSA indexer backend interface to declare block-cache requests.
lmdeploy/pytorch/backends/indexer.py	Updates V4 indexer backend interface to accept block_caches mapping.
lmdeploy/pytorch/backends/dlinfer/op_backend.py	Adds cache-backend provider hook for dlinfer backend.
lmdeploy/pytorch/backends/dlinfer/cache.py	Adds dlinfer cache backend selecting contiguous layouts.
lmdeploy/pytorch/backends/default/op_backend.py	Adds default cache-backend provider hook and removes old CacheBlockCopy op builder.
lmdeploy/pytorch/backends/default/cache.py	Adds default cache backend selecting layouts and torch fallback block-copy primitive.
lmdeploy/pytorch/backends/default/cache_block_copy.py	Removes legacy op-based cache block copy builder/impl.
lmdeploy/pytorch/backends/cuda/v4_indexer.py	Updates CUDA V4 indexer to read caches from block_caches mapping using standardized names.
lmdeploy/pytorch/backends/cuda/v4_compressor.py	Adds cache request declarations for V4 compressor/indexer variants and updates write interface to named caches.
lmdeploy/pytorch/backends/cuda/op_backend.py	Adds CUDA cache-backend provider hook and removes old CacheBlockCopy op builder.
lmdeploy/pytorch/backends/cuda/nsa.py	Implements DSA indexer cache request emission (requires aligned logical/kernel geometry).
lmdeploy/pytorch/backends/cuda/cache.py	Adds CUDA cache backend providing Triton-based block copy over allocation pools.
lmdeploy/pytorch/backends/cuda/cache_block_copy.py	Removes legacy op-based CUDA cache block copy builder/impl.
lmdeploy/pytorch/backends/cuda/attention/v4.py	Updates V4 attention executor to accept window state + named block caches and resolves compressed KV cache by standardized name.
lmdeploy/pytorch/backends/compressor.py	Extends compressor backend interface to declare block-cache requests and accept named caches for writes.
lmdeploy/pytorch/backends/cache.py	Introduces backend cache abstractions (CacheBackend, BlockCacheCopy) and allocation validation.
lmdeploy/pytorch/backends/cache_block_copy.py	Removes legacy CacheBlockCopy op abstraction.
lmdeploy/pytorch/backends/base.py	Adds get_cache_backend requirement to backends and removes OpType.CacheBlockCopy.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.