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.