[Runtime] Add PagedAttentionKVCache checkpoint primitives#20035
[Runtime] Add PagedAttentionKVCache checkpoint primitives#20035akaashrp wants to merge 3 commits into
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces checkpointing support for the PagedAttentionKVCache in the TVM runtime VM, enabling exporting and importing of page groups, retrieving checkpoint metadata, and verifying cache layouts via stable hashing. The changes are accompanied by comprehensive unit tests. The review feedback highlights critical improvements to ensure robustness and portability: explicitly validating that qk_head_dim_ equals v_head_dim_ to prevent undefined behavior, and making the layout descriptor serialization locale-independent by imbuing the stream with std::locale::classic() and setting the stream precision, along with including the <locale> header.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Add checkpoint, export, and import primitives for
PagedAttentionKVCache. This is intended to support resuming generation after a crash in WebLLM without requiring model ABI changes. The new runtime APIs expose enough cache metadata and page data for an external runtime to persist KV state, restore the cache after reload, and validate that the restored bytes match the current cache layout.