[Runtime] Add PagedAttentionKVCache checkpoint primitives

[akaashrp]

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.

[gemini-code-assist]

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.

[gemini-code-assist]

The checkpointing export and import logic assumes that the query-key head dimension (qk_head_dim_) is equal to the value head dimension (v_head_dim_) when constructing and verifying page tensor shapes (e.g., in MakePageGroupMetadata and CheckPageGroupTensor). To prevent undefined behavior or crashes if they differ, explicitly validate that qk_head_dim_ == v_head_dim_ in CheckCheckpointLayoutSupported().

    TVM_FFI_ICHECK_EQ(qk_head_dim_, v_head_dim_)
        << "PagedAttentionKVCache checkpointing requires qk_head_dim to equal v_head_dim.";
    TVM_FFI_ICHECK(!support_sliding_window_ && !support_layer_sliding_window_)
        << "PagedAttentionKVCache checkpointing does not support sliding-window cache layouts.";
    TVM_FFI_ICHECK(!f_transfer_kv_.has_value() && !f_transfer_kv_page_to_page_.has_value())
        << "PagedAttentionKVCache checkpointing does not support KV transfer/disaggregation.";

[gemini-code-assist]

The floating-point formatting of rotary_scale_ and rotary_theta_ in GetLayoutDescriptor() is locale-dependent and subject to default stream precision. This can lead to unstable layout hashes across different environments or locales (e.g., using commas instead of dots as decimal separators). Imbue the stream with std::locale::classic() and set the precision to std::numeric_limits<double>::max_digits10 to ensure stable, locale-independent hashing.

Suggested change

	std::string GetLayoutDescriptor() const {
	std::ostringstream os;
	os << "cacheType=" << kPagedKVCacheCheckpointRuntime << ";";
	os << "pageSize=" << page_size_ << ";";
	os << "numLayers=" << num_layers_ << ";";
	os << "layerBegin=" << layer_id_begin_offset_ << ";";
	os << "layerEnd=" << layer_id_end_offset_ << ";";
	os << "numQOHeads=" << num_qo_heads_ << ";";
	os << "numKVHeads=" << num_kv_heads_ << ";";
	os << "qkHeadDim=" << qk_head_dim_ << ";";
	os << "vHeadDim=" << v_head_dim_ << ";";
	os << "numTotalPages=" << num_total_pages_ << ";";
	os << "prefillChunkSize=" << prefill_chunk_size_ << ";";
	os << "dtype=" << std::string(ffi::DLDataTypeToString(kv_dtype_)) << ";";
	os << "ropeMode=" << RoPEModeToString(rope_mode_) << ";";
	os << "rotaryScale=" << rotary_scale_ << ";";
	os << "rotaryTheta=" << rotary_theta_ << ";";
	std::string GetLayoutDescriptor() const {
	std::ostringstream os;
	os.imbue(std::locale::classic());
	os << std::setprecision(std::numeric_limits<double>::max_digits10);
	os << "cacheType=" << kPagedKVCacheCheckpointRuntime << ";";
	os << "pageSize=" << page_size_ << ";";
	os << "numLayers=" << num_layers_ << ";";
	os << "layerBegin=" << layer_id_begin_offset_ << ";";
	os << "layerEnd=" << layer_id_end_offset_ << ";";
	os << "numQOHeads=" << num_qo_heads_ << ";";
	os << "numKVHeads=" << num_kv_heads_ << ";";
	os << "qkHeadDim=" << qk_head_dim_ << ";";
	os << "vHeadDim=" << v_head_dim_ << ";";
	os << "numTotalPages=" << num_total_pages_ << ";";
	os << "prefillChunkSize=" << prefill_chunk_size_ << ";";
	os << "dtype=" << std::string(ffi::DLDataTypeToString(kv_dtype_)) << ";";
	os << "ropeMode=" << RoPEModeToString(rope_mode_) << ";";
	os << "rotaryScale=" << rotary_scale_ << ";";
	os << "rotaryTheta=" << rotary_theta_ << ";";

[gemini-code-assist]

Include <locale> to support std::locale::classic() when formatting the layout descriptor stream.

Suggested change

	#include <iomanip>
	#include <limits>
	#include <numeric>
	#include <sstream>
	#include <iomanip>
	#include <limits>
	#include <locale>
	#include <numeric>
	#include <sstream>