[Runtime] Add PagedAttentionKVCache checkpoint primitives

src/runtime/vm/kv_state.cc

@@ -70,6 +70,18 @@ TVM_FFI_STATIC_INIT_BLOCK() {
                   &AttentionKVCacheObj::GetNumAvailablePages)
       .def_method("vm.builtin.attention_kv_cache_get_total_sequence_length",
                   &AttentionKVCacheObj::GetTotalSequenceLength)
+      .def_method("vm.builtin.attention_kv_cache_get_checkpoint_metadata",
+                  &AttentionKVCacheObj::GetCheckpointMetadata)
+      .def_method("vm.builtin.attention_kv_cache_get_layout_hash",
+                  &AttentionKVCacheObj::GetLayoutHash)
+      .def_method("vm.builtin.attention_kv_cache_export_page_group",
+                  &AttentionKVCacheObj::ExportPageGroup)
+      .def_method("vm.builtin.attention_kv_cache_prepare_import",
+                  &AttentionKVCacheObj::PrepareImport)
+      .def_method("vm.builtin.attention_kv_cache_import_page_group",
+                  &AttentionKVCacheObj::ImportPageGroup)
+      .def_method("vm.builtin.attention_kv_cache_get_sequence_length",
+                  &AttentionKVCacheObj::GetSequenceLength)
       .def_method("vm.builtin.attention_kv_cache_get_query_positions",
                   &AttentionKVCacheObj::GetQueryPositions)
       .def_method("vm.builtin.attention_kv_cache_debug_get_kv", &AttentionKVCacheObj::DebugGetKV)

src/runtime/vm/kv_state.h

@@ -23,6 +23,7 @@
 #include <tvm/ffi/error.h>
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/optional.h>
+#include <tvm/ffi/string.h>
 #include <tvm/runtime/device_api.h>
 #include <tvm/runtime/tensor.h>
 
@@ -135,6 +136,49 @@ class AttentionKVCacheObj : public KVStateObj {
   /*! \brief Get the current total sequence length in the KV cache. */
   virtual int32_t GetTotalSequenceLength() const = 0;
 
+  /*!
+   * \brief Get checkpoint metadata for a sequence.
+   * \param seq_id The id of the sequence whose checkpoint metadata is requested.
+   * \return JSON string describing runtime layout and logical pages.
+   */
+  virtual ffi::String GetCheckpointMetadata(int64_t seq_id) const = 0;
+
+  /*!
+   * \brief Get a stable hash over runtime layout-defining fields.
+   * \return Hex string hash of the checkpoint layout metadata.
+   */
+  virtual ffi::String GetLayoutHash() const = 0;
+
+  /*!
+   * \brief Export a checkpoint page group for a sequence.
+   * \param seq_id The id of the sequence whose page group is exported.
+   * \param group_id The checkpoint group id to export.
+   * \param dst The destination tensor for the exported page group.
+   */
+  virtual void ExportPageGroup(int64_t seq_id, int64_t group_id, Tensor dst) = 0;
+
+  /*!
+   * \brief Prepare sequence state and page tables for checkpoint import.
+   * \param seq_id The id of the sequence being imported.
+   * \param metadata_json The checkpoint metadata JSON string.
+   */
+  virtual void PrepareImport(int64_t seq_id, ffi::String metadata_json) = 0;
+
+  /*!
+   * \brief Import a checkpoint page group into a prepared sequence.
+   * \param seq_id The id of the sequence being imported.
+   * \param group_id The checkpoint group id to import.
+   * \param src The source tensor containing the page group.
+   */
+  virtual void ImportPageGroup(int64_t seq_id, int64_t group_id, Tensor src) = 0;
+
+  /*!
+   * \brief Get the sequence length for a sequence in the KV cache.
+   * \param seq_id The id of the sequence whose length is requested.
+   * \return The sequence length.
+   */
+  virtual int32_t GetSequenceLength(int64_t seq_id) const = 0;
+
   /************** Sequence Management **************/
 
   /*!