This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 62408ab065 [Runtime] Add PagedAttentionKVCache checkpoint primitives
(#20035)
62408ab065 is described below
commit 62408ab065fcd8b8c993a479271610a1419f093b
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Tue Aug 25 15:50:56 2026 -0700
[Runtime] Add PagedAttentionKVCache checkpoint primitives (#20035)
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.
---
src/runtime/vm/kv_state.cc | 13 +
src/runtime/vm/kv_state.h | 50 ++
src/runtime/vm/paged_kv_cache.cc | 698 ++++++++++++++++++++-
...runtime_builtin_paged_attention_kv_cache_cpu.py | 140 +++++
...me_builtin_paged_attention_kv_cache_metadata.py | 417 ++++++++++++
5 files changed, 1315 insertions(+), 3 deletions(-)
diff --git a/src/runtime/vm/kv_state.cc b/src/runtime/vm/kv_state.cc
index 2d93563497..3946456c36 100644
--- a/src/runtime/vm/kv_state.cc
+++ b/src/runtime/vm/kv_state.cc
@@ -70,6 +70,19 @@ 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_finish_import",
&AttentionKVCacheObj::FinishImport)
+ .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)
diff --git a/src/runtime/vm/kv_state.h b/src/runtime/vm/kv_state.h
index 4e1694c899..725af90a42 100644
--- a/src/runtime/vm/kv_state.h
+++ b/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,55 @@ 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 Finish a checkpoint import after all page groups have been
restored.
+ * \param seq_id The id of the sequence being imported.
+ */
+ virtual void FinishImport(int64_t seq_id) = 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 **************/
/*!
diff --git a/src/runtime/vm/paged_kv_cache.cc b/src/runtime/vm/paged_kv_cache.cc
index cd0722f865..387424f5e7 100644
--- a/src/runtime/vm/paged_kv_cache.cc
+++ b/src/runtime/vm/paged_kv_cache.cc
@@ -22,6 +22,7 @@
*/
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/error.h>
+#include <tvm/ffi/extra/json.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
@@ -31,7 +32,11 @@
#include <tvm/support/cuda/nvtx.h>
#include <algorithm>
+#include <iomanip>
+#include <limits>
+#include <locale>
#include <numeric>
+#include <sstream>
#include <unordered_map>
#include <utility>
#include <vector>
@@ -44,6 +49,57 @@ namespace tvm {
namespace runtime {
namespace vm {
+namespace {
+
+constexpr const char* kPagedKVCacheCheckpointRuntime =
"relax.vm.PagedAttentionKVCache";
+constexpr int64_t kPagedKVCacheCheckpointFormatVersion = 1;
+
+const char* AttnKindToString(AttnKind attn_kind) {
+ switch (attn_kind) {
+ case AttnKind::kMHA:
+ return "mha";
+ case AttnKind::kMLA:
+ return "mla";
+ case AttnKind::kLinearAttn:
+ return "linear";
+ case AttnKind::kMHASliding:
+ return "mha_sliding";
+ }
+ TVM_FFI_ICHECK(false) << "Unknown attention kind: " <<
static_cast<int>(attn_kind);
+ return "unknown";
+}
+
+const char* RoPEModeToString(RoPEMode rope_mode) {
+ switch (rope_mode) {
+ case RoPEMode::kNone:
+ return "none";
+ case RoPEMode::kNormal:
+ return "normal";
+ case RoPEMode::kInline:
+ return "inline";
+ }
+ TVM_FFI_ICHECK(false) << "Unknown RoPE mode: " <<
static_cast<int>(rope_mode);
+ return "unknown";
+}
+
+std::string Uint64ToHex(uint64_t value) {
+ std::ostringstream os;
+ os.imbue(std::locale::classic());
+ os << std::hex << std::setw(16) << std::setfill('0') << value;
+ return os.str();
+}
+
+uint64_t FNV1a64(const std::string& value) {
+ uint64_t hash = 14695981039346656037ULL;
+ for (unsigned char c : value) {
+ hash ^= c;
+ hash *= 1099511628211ULL;
+ }
+ return hash;
+}
+
+} // namespace
+
//-------------------------------------------
// We keep the implementation private as
// they may subject to future changes.
@@ -95,6 +151,8 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj {
* For layers that use multi-head attention, this field is overriden by
qk_head_dim.
*/
const int64_t v_head_dim_;
+ /*! \brief The number of sequences reserved in the KV cache. */
+ const int64_t reserved_num_seqs_;
/*! \brief The number of total pages allocated in KV cache. */
const int64_t num_total_pages_;
/*! \brief The maximum total sequence length in a prefill. */
@@ -133,12 +191,18 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
* Along on the "2" dimension, index 0 stands for K and 1 stands for V.
*/
std::vector<Tensor> pages_;
+ /*! \brief A reusable host-side zero page for deterministic checkpoint
padding. */
+ Tensor checkpoint_zero_page_host_;
/*! \brief The whole KV cache allocated by NVSHMEM*/
Tensor nvshmem_pages_;
/*! \brief The list of ids of released pages for page reuse. */
std::vector<int32_t> free_page_ids_;
/*! \brief The mapping from sequence ids to sequences. */
std::unordered_map<int64_t, Sequence> seq_map_;
+ /*! \brief Whether a checkpoint import is waiting for all page groups. */
+ bool checkpoint_import_in_progress_ = false;
+ /*! \brief The page groups restored by the current checkpoint import. */
+ std::vector<bool> checkpoint_imported_groups_;
/********************* Sequence Block Structures *********************/
@@ -161,8 +225,6 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj
{
bool dirty_aux_data_device_ = false;
/*! \brief The batch size of the current round of forwarding. */
int64_t cur_batch_size_;
- /*! \brief The number of sequences reserved in the KV cache. */
- int64_t reserved_num_seqs_;
/*! \brief The ids of the sequences in the current round of forwarding. */
ffi::Shape cur_seq_ids_;
/*! \brief The append lengths of the sequences in the current round of
forwarding. */
@@ -330,6 +392,7 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj
{
num_kv_heads_(num_kv_heads),
qk_head_dim_(qk_head_dim),
v_head_dim_(v_head_dim),
+ reserved_num_seqs_(reserved_num_seqs),
num_total_pages_(num_total_pages),
prefill_chunk_size_(prefill_chunk_size),
support_sliding_window_(std::find(attn_kinds.begin(), attn_kinds.end(),
@@ -346,7 +409,6 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj
{
rotary_theta_(rotary_theta),
rope_ext_factors_(std::move(rope_ext_factors)),
kv_dtype_(dtype),
- reserved_num_seqs_(reserved_num_seqs),
f_transpose_append_mha_(std::move(f_transpose_append_mha)),
f_transpose_append_mla_(std::move(f_transpose_append_mla)),
f_compact_copy_(std::move(f_compact_copy)),
@@ -560,6 +622,8 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj
{
/*! \brief Reset the KV cache. */
void Clear() final {
seq_map_.clear();
+ checkpoint_import_in_progress_ = false;
+ checkpoint_imported_groups_.clear();
free_page_ids_.clear();
for (int64_t page_id = num_total_pages_ - 1; page_id >= 0; --page_id) {
free_page_ids_.push_back(page_id);
@@ -879,10 +943,161 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
return total_seq_len;
}
+ ffi::String GetCheckpointMetadata(int64_t seq_id) const final {
+ CheckCheckpointSequenceSupported(seq_id);
+ const Sequence& seq = seq_map_.at(seq_id);
+
+ namespace json = tvm::ffi::json;
+ json::Object metadata = MakeLayoutMetadata();
+ metadata.Set("layout_hash", GetLayoutHash());
+ metadata.Set("seq_id", seq_id);
+ metadata.Set("seq_length", static_cast<int64_t>(seq.seq_length));
+ metadata.Set("logical_pages", MakeLogicalPageMetadata(seq));
+ metadata.Set("groups", MakePageGroupMetadata(seq));
+ return json::Stringify(metadata);
+ }
+
+ ffi::String GetLayoutHash() const final {
+ CheckCheckpointLayoutSupported();
+ return ffi::String(Uint64ToHex(FNV1a64(GetLayoutDescriptor())));
+ }
+
+ void ExportPageGroup(int64_t seq_id, int64_t group_id, Tensor dst) final {
+ CheckCheckpointSequenceSupported(seq_id);
+ const Sequence& seq = seq_map_.at(seq_id);
+ TVM_FFI_ICHECK_GE(group_id, 0)
+ << "PagedAttentionKVCache checkpoint export got invalid group id " <<
group_id << ".";
+ TVM_FFI_ICHECK_LT(group_id, num_layers_)
+ << "PagedAttentionKVCache checkpoint export got invalid group id " <<
group_id
+ << ", but only " << num_layers_ << " groups are available.";
+
+ int64_t num_logical_pages = GetNumLogicalPages(seq);
+ CheckExportPageGroupTensor(dst, num_logical_pages);
+ std::vector<int32_t> page_ids = GetCheckpointPageIds(seq, "checkpoint
export");
+
+ if (copy_stream_ != nullptr) {
+ DeviceAPI::Get(device_)->SyncStreamFromTo(device_, copy_stream_,
compute_stream_);
+ }
+
+ Tensor layer_pages = pages_[group_id];
+ int64_t num_full_pages = seq.seq_length / page_size_;
+ int64_t partial_page_length = seq.seq_length % page_size_;
+ if (partial_page_length != 0) {
+ TVM_FFI_ICHECK_LT(num_full_pages, num_logical_pages);
+ ExportPartialCheckpointPage(layer_pages, page_ids[num_full_pages], dst,
num_full_pages,
+ partial_page_length);
+ }
+
+ int64_t logical_page_index = 0;
+ while (logical_page_index < num_full_pages) {
+ int64_t run_length = 1;
+ while (logical_page_index + run_length < num_full_pages &&
+ page_ids[logical_page_index + run_length] ==
+ page_ids[logical_page_index] + run_length) {
+ ++run_length;
+ }
+ CopyCheckpointPageRun(layer_pages, page_ids[logical_page_index], dst,
logical_page_index,
+ run_length);
+ logical_page_index += run_length;
+ }
+ if (partial_page_length != 0) {
+ ++logical_page_index;
+ }
+ TVM_FFI_ICHECK_EQ(logical_page_index, num_logical_pages);
+ }
+
+ void PrepareImport(int64_t seq_id, ffi::String metadata_json) final {
+ CheckCheckpointLayoutSupported();
+ TVM_FFI_ICHECK_EQ(seq_id, 0)
+ << "PagedAttentionKVCache checkpoint import only supports sequence id
0, got " << seq_id
+ << ".";
+ ffi::json::Object metadata = ParseCheckpointMetadata(metadata_json);
+ int64_t seq_length = CheckCheckpointImportMetadata(seq_id, metadata);
+ int64_t num_logical_pages = GetExpectedNumLogicalPages(seq_length);
+
+ Clear();
+ int32_t block_idx = GetFreeBlock();
+ Block& block = global_block_pool_[block_idx];
+ block.start_pos = 0;
+ block.seq_length = static_cast<int32_t>(seq_length);
+ for (int64_t page_index = 0; page_index < num_logical_pages; ++page_index)
{
+ block.page_ids.push_back(GetFreePage());
+ }
+ seq_map_.insert({seq_id, Sequence(&global_block_pool_, block_idx)});
+ checkpoint_import_in_progress_ = true;
+ checkpoint_imported_groups_.assign(num_layers_, false);
+ dirty_aux_data_device_ = true;
+ }
+
+ void ImportPageGroup(int64_t seq_id, int64_t group_id, Tensor src) final {
+ CheckCheckpointLayoutSupported();
+ TVM_FFI_ICHECK(checkpoint_import_in_progress_)
+ << "PagedAttentionKVCache checkpoint import has not been prepared.";
+ CheckCheckpointSequenceState(seq_id);
+ const Sequence& seq = seq_map_.at(seq_id);
+ TVM_FFI_ICHECK_GE(group_id, 0)
+ << "PagedAttentionKVCache checkpoint import got invalid group id " <<
group_id << ".";
+ TVM_FFI_ICHECK_LT(group_id, num_layers_)
+ << "PagedAttentionKVCache checkpoint import got invalid group id " <<
group_id
+ << ", but only " << num_layers_ << " groups are available.";
+ TVM_FFI_ICHECK(!checkpoint_imported_groups_[group_id])
+ << "PagedAttentionKVCache checkpoint group " << group_id << " was
already imported.";
+
+ int64_t num_logical_pages = GetNumLogicalPages(seq);
+ CheckImportPageGroupTensor(src, num_logical_pages);
+ std::vector<int32_t> page_ids = GetCheckpointPageIds(seq, "checkpoint
import");
+
+ if (copy_stream_ != nullptr) {
+ DeviceAPI::Get(device_)->SyncStreamFromTo(device_, copy_stream_,
compute_stream_);
+ }
+
+ Tensor layer_pages = pages_[group_id];
+ int64_t logical_page_index = 0;
+ while (logical_page_index < num_logical_pages) {
+ int64_t run_length = 1;
+ while (logical_page_index + run_length < num_logical_pages &&
+ page_ids[logical_page_index + run_length] ==
+ page_ids[logical_page_index] + run_length) {
+ ++run_length;
+ }
+ CopyCheckpointPageRun(src, logical_page_index, layer_pages,
page_ids[logical_page_index],
+ run_length);
+ logical_page_index += run_length;
+ }
+ TVM_FFI_ICHECK_EQ(logical_page_index, num_logical_pages);
+ checkpoint_imported_groups_[group_id] = true;
+ }
+
+ void FinishImport(int64_t seq_id) final {
+ CheckCheckpointLayoutSupported();
+ TVM_FFI_ICHECK(checkpoint_import_in_progress_)
+ << "PagedAttentionKVCache checkpoint import has not been prepared.";
+ CheckCheckpointSequenceState(seq_id);
+ for (int64_t group_id = 0; group_id < num_layers_; ++group_id) {
+ TVM_FFI_ICHECK(checkpoint_imported_groups_[group_id])
+ << "PagedAttentionKVCache checkpoint import is missing group " <<
group_id << ".";
+ }
+ checkpoint_import_in_progress_ = false;
+ checkpoint_imported_groups_.clear();
+ }
+
+ int32_t GetSequenceLength(int64_t seq_id) const final {
+ TVM_FFI_ICHECK(!checkpoint_import_in_progress_)
+ << "PagedAttentionKVCache sequence length is unavailable until
checkpoint import is "
+ "finished.";
+ auto it = seq_map_.find(seq_id);
+ TVM_FFI_ICHECK(it != seq_map_.end())
+ << "The sequence \"" << seq_id << "\" cannot be found in KV cache.";
+ return it->second.seq_length;
+ }
+
/************** Attention **************/
void BeginForward(const ffi::Shape& seq_ids, const ffi::Shape&
append_lengths,
const ffi::Optional<ffi::Shape>&
opt_token_tree_parent_ptr) final {
+ TVM_FFI_ICHECK(!checkpoint_import_in_progress_)
+ << "PagedAttentionKVCache cannot begin a forward pass before
checkpoint import is "
+ "finished.";
// Note: MLA does not supported tree attention for now.
if (attn_kinds_[0] == AttnKind::kMLA) {
TVM_FFI_ICHECK(!opt_token_tree_parent_ptr.has_value())
@@ -1773,6 +1988,483 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
AttentionKVCacheObj);
private:
+ void CheckCheckpointLayoutSupported() const {
+ for (int64_t layer_id = layer_id_begin_offset_; layer_id <
layer_id_end_offset_; ++layer_id) {
+ AttnKind attn_kind = attn_kinds_[layer_id];
+ TVM_FFI_ICHECK(attn_kind == AttnKind::kMHA)
+ << "PagedAttentionKVCache checkpointing only supports full-context
MHA/GQA layers.";
+ }
+ 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(!rope_ext_factors_.has_value())
+ << "PagedAttentionKVCache checkpointing does not support RoPE
extension factors.";
+ TVM_FFI_ICHECK(!f_transfer_kv_.has_value() &&
!f_transfer_kv_page_to_page_.has_value())
+ << "PagedAttentionKVCache checkpointing does not support KV
transfer/disaggregation.";
+ }
+
+ void CheckCheckpointSequenceSupported(int64_t seq_id) const {
+ CheckCheckpointLayoutSupported();
+ TVM_FFI_ICHECK(!checkpoint_import_in_progress_)
+ << "PagedAttentionKVCache checkpoint import must be finished before
exporting a "
+ "checkpoint.";
+ CheckCheckpointSequenceState(seq_id);
+ }
+
+ void CheckCheckpointSequenceState(int64_t seq_id) const {
+ TVM_FFI_ICHECK_EQ(seq_id, 0)
+ << "PagedAttentionKVCache checkpointing only supports sequence id 0,
got " << seq_id << ".";
+ auto it = seq_map_.find(seq_id);
+ TVM_FFI_ICHECK(it != seq_map_.end())
+ << "The sequence \"" << seq_id << "\" cannot be found in KV cache.";
+ const Sequence& seq = it->second;
+ TVM_FFI_ICHECK(seq.accepted_indices_committed && seq.is_chain)
+ << "PagedAttentionKVCache checkpointing requires committed token-chain
state.";
+ TVM_FFI_ICHECK_EQ(seq.sliding_window_size, -1)
+ << "PagedAttentionKVCache checkpointing does not support sequences
with sliding window.";
+ }
+
+ int64_t GetExpectedNumLogicalPages(int64_t seq_length) const {
+ return (seq_length + page_size_ - 1) / page_size_;
+ }
+
+ ffi::json::Object ParseCheckpointMetadata(const ffi::String& metadata_json)
const {
+ ffi::String error_msg;
+ ffi::json::Value json_info = ffi::json::Parse(metadata_json, &error_msg);
+ TVM_FFI_ICHECK(error_msg.empty())
+ << "Failed to parse PagedAttentionKVCache checkpoint metadata JSON: "
<< error_msg << ".";
+ TVM_FFI_ICHECK(json_info.as<ffi::json::Object>())
+ << "PagedAttentionKVCache checkpoint metadata should be a JSON
object.";
+ return json_info.cast<ffi::json::Object>();
+ }
+
+ ffi::json::Value GetJSONField(const ffi::json::Object& object, const char*
field,
+ const char* context) const {
+ auto it = object.find(field);
+ TVM_FFI_ICHECK(it != object.end()) << context << " missing field \"" <<
field << "\".";
+ return (*it).second;
+ }
+
+ int64_t GetJSONIntegerField(const ffi::json::Object& object, const char*
field,
+ const char* context) const {
+ return GetJSONField(object, field, context).cast<int64_t>();
+ }
+
+ double GetJSONNumberField(const ffi::json::Object& object, const char* field,
+ const char* context) const {
+ return GetJSONField(object, field, context).cast<double>();
+ }
+
+ bool GetJSONBoolField(const ffi::json::Object& object, const char* field,
+ const char* context) const {
+ return GetJSONField(object, field, context).cast<bool>();
+ }
+
+ std::string GetJSONStringField(const ffi::json::Object& object, const char*
field,
+ const char* context) const {
+ return std::string(GetJSONField(object, field,
context).cast<ffi::String>());
+ }
+
+ ffi::json::Array GetJSONArrayField(const ffi::json::Object& object, const
char* field,
+ const char* context) const {
+ return GetJSONField(object, field, context).cast<ffi::json::Array>();
+ }
+
+ void CheckJSONIntegerField(const ffi::json::Object& object, const char*
field, int64_t expected,
+ const char* context) const {
+ int64_t value = GetJSONIntegerField(object, field, context);
+ TVM_FFI_ICHECK_EQ(value, expected)
+ << context << " field \"" << field << "\" mismatch: expected " <<
expected << ", got "
+ << value << ".";
+ }
+
+ void CheckJSONStringField(const ffi::json::Object& object, const char* field,
+ const std::string& expected, const char* context)
const {
+ std::string value = GetJSONStringField(object, field, context);
+ TVM_FFI_ICHECK_EQ(value, expected)
+ << context << " field \"" << field << "\" mismatch: expected " <<
expected << ", got "
+ << value << ".";
+ }
+
+ void CheckJSONNumberField(const ffi::json::Object& object, const char*
field, double expected,
+ const char* context) const {
+ double value = GetJSONNumberField(object, field, context);
+ TVM_FFI_ICHECK_EQ(value, expected)
+ << context << " field \"" << field << "\" mismatch: expected " <<
expected << ", got "
+ << value << ".";
+ }
+
+ void CheckJSONBoolField(const ffi::json::Object& object, const char* field,
bool expected,
+ const char* context) const {
+ bool value = GetJSONBoolField(object, field, context);
+ TVM_FFI_ICHECK_EQ(value, expected)
+ << context << " field \"" << field << "\" mismatch: expected " <<
expected << ", got "
+ << value << ".";
+ }
+
+ void CheckCheckpointImportLayout(const ffi::json::Object& metadata) const {
+ static constexpr const char* context = "PagedAttentionKVCache checkpoint
import metadata";
+ CheckJSONIntegerField(metadata, "format_version",
kPagedKVCacheCheckpointFormatVersion,
+ context);
+ CheckJSONStringField(metadata, "cache_type",
kPagedKVCacheCheckpointRuntime, context);
+ CheckJSONIntegerField(metadata, "page_size", page_size_, context);
+ CheckJSONIntegerField(metadata, "num_layers", num_layers_, context);
+ CheckJSONIntegerField(metadata, "layer_begin", layer_id_begin_offset_,
context);
+ CheckJSONIntegerField(metadata, "layer_end", layer_id_end_offset_,
context);
+ CheckJSONIntegerField(metadata, "num_qo_heads", num_qo_heads_, context);
+ CheckJSONIntegerField(metadata, "num_kv_heads", num_kv_heads_, context);
+ CheckJSONIntegerField(metadata, "qk_head_dim", qk_head_dim_, context);
+ CheckJSONIntegerField(metadata, "v_head_dim", v_head_dim_, context);
+ CheckJSONStringField(metadata, "dtype",
std::string(ffi::DLDataTypeToString(kv_dtype_)),
+ context);
+ CheckJSONStringField(metadata, "rope_mode", RoPEModeToString(rope_mode_),
context);
+ CheckJSONNumberField(metadata, "rotary_scale", rotary_scale_, context);
+ CheckJSONNumberField(metadata, "rotary_theta", rotary_theta_, context);
+ CheckJSONBoolField(metadata, "has_rope_ext_factors",
rope_ext_factors_.has_value(), context);
+ CheckJSONBoolField(metadata, "support_sliding_window",
support_sliding_window_, context);
+ CheckJSONBoolField(metadata, "support_layer_sliding_window",
support_layer_sliding_window_,
+ context);
+ CheckJSONStringField(metadata, "page_group_layout",
+
"1,num_logical_pages,2,num_kv_heads,page_size,qk_head_dim", context);
+ TVM_FFI_ICHECK_GT(GetJSONIntegerField(metadata, "reserved_num_seqs",
context), 0)
+ << context << " field \"reserved_num_seqs\" must be positive.";
+ TVM_FFI_ICHECK_GT(GetJSONIntegerField(metadata, "num_total_pages",
context), 0)
+ << context << " field \"num_total_pages\" must be positive.";
+ TVM_FFI_ICHECK_GT(GetJSONIntegerField(metadata, "prefill_chunk_size",
context), 0)
+ << context << " field \"prefill_chunk_size\" must be positive.";
+ ffi::String expected_layout_hash = GetLayoutHash();
+ std::string layout_hash = GetJSONStringField(metadata, "layout_hash",
context);
+ TVM_FFI_ICHECK_EQ(layout_hash, std::string(expected_layout_hash))
+ << "PagedAttentionKVCache checkpoint import layout hash mismatch:
expected "
+ << expected_layout_hash << ", got " << layout_hash << ".";
+
+ ffi::json::Array attn_kinds = GetJSONArrayField(metadata, "attn_kinds",
context);
+ TVM_FFI_ICHECK_EQ(attn_kinds.size(), num_layers_)
+ << context << " field \"attn_kinds\" size mismatch.";
+ for (int64_t local_layer = 0; local_layer < num_layers_; ++local_layer) {
+ std::string attn_kind =
std::string(attn_kinds[local_layer].cast<ffi::String>());
+ std::string expected =
AttnKindToString(attn_kinds_[layer_id_begin_offset_ + local_layer]);
+ TVM_FFI_ICHECK_EQ(attn_kind, expected)
+ << context << " field \"attn_kinds\" mismatch at local layer " <<
local_layer
+ << ": expected " << expected << ", got " << attn_kind << ".";
+ }
+ }
+
+ void CheckCheckpointImportLogicalPages(const ffi::json::Object& metadata,
+ int64_t seq_length) const {
+ static constexpr const char* context = "PagedAttentionKVCache checkpoint
import metadata";
+ int64_t num_logical_pages = GetExpectedNumLogicalPages(seq_length);
+ ffi::json::Array logical_pages = GetJSONArrayField(metadata,
"logical_pages", context);
+ TVM_FFI_ICHECK_EQ(static_cast<int64_t>(logical_pages.size()),
num_logical_pages)
+ << "PagedAttentionKVCache checkpoint import sequence length mismatch:
seq_length "
+ << seq_length << " requires " << num_logical_pages << " logical pages,
but metadata has "
+ << logical_pages.size() << ".";
+
+ int64_t expected_start = 0;
+ for (int64_t i = 0; i < num_logical_pages; ++i) {
+ ffi::json::Object page = logical_pages[i].cast<ffi::json::Object>();
+ int64_t expected_length = std::min<int64_t>(page_size_, seq_length -
expected_start);
+ CheckJSONIntegerField(page, "logical_page_index", i, context);
+ CheckJSONIntegerField(page, "start_pos", expected_start, context);
+ CheckJSONIntegerField(page, "length", expected_length, context);
+ expected_start += expected_length;
+ }
+ TVM_FFI_ICHECK_EQ(expected_start, seq_length)
+ << "PagedAttentionKVCache checkpoint import logical pages do not cover
seq_length "
+ << seq_length << ".";
+ }
+
+ void CheckPageGroupMetadataShape(const ffi::json::Object& group, int64_t
num_logical_pages,
+ const char* context) const {
+ ffi::json::Array shape = GetJSONArrayField(group, "shape", context);
+ std::vector<int64_t> expected_shape = {1, num_logical_pages, 2,
num_kv_heads_,
+ page_size_, qk_head_dim_};
+ TVM_FFI_ICHECK_EQ(shape.size(), expected_shape.size())
+ << context << " field \"shape\" rank mismatch.";
+ for (int64_t i = 0; i < static_cast<int64_t>(expected_shape.size()); ++i) {
+ int64_t dim = shape[i].cast<int64_t>();
+ TVM_FFI_ICHECK_EQ(dim, expected_shape[i])
+ << context << " field \"shape\" mismatch at dim " << i << ":
expected "
+ << expected_shape[i] << ", got " << dim << ".";
+ }
+ }
+
+ void CheckCheckpointImportGroups(const ffi::json::Object& metadata,
+ int64_t num_logical_pages) const {
+ static constexpr const char* context = "PagedAttentionKVCache checkpoint
import group metadata";
+ ffi::json::Array groups = GetJSONArrayField(metadata, "groups", context);
+ TVM_FFI_ICHECK_EQ(static_cast<int64_t>(groups.size()), num_layers_)
+ << context << " size mismatch: expected " << num_layers_ << ", got "
<< groups.size()
+ << ".";
+ for (int64_t local_layer = 0; local_layer < num_layers_; ++local_layer) {
+ ffi::json::Object group = groups[local_layer].cast<ffi::json::Object>();
+ CheckJSONIntegerField(group, "group_index", local_layer, context);
+ CheckJSONIntegerField(group, "layer_begin", layer_id_begin_offset_ +
local_layer, context);
+ CheckJSONIntegerField(group, "layer_end", layer_id_begin_offset_ +
local_layer + 1, context);
+ CheckJSONIntegerField(group, "num_logical_pages", num_logical_pages,
context);
+ CheckJSONStringField(group, "dtype",
std::string(ffi::DLDataTypeToString(kv_dtype_)),
+ context);
+ CheckPageGroupMetadataShape(group, num_logical_pages, context);
+ CheckJSONIntegerField(group, "nbytes",
GetCheckpointGroupNBytes(num_logical_pages), context);
+ }
+ }
+
+ int64_t CheckCheckpointImportMetadata(int64_t seq_id, const
ffi::json::Object& metadata) const {
+ CheckCheckpointImportLayout(metadata);
+ CheckJSONIntegerField(metadata, "seq_id", seq_id,
+ "PagedAttentionKVCache checkpoint import metadata");
+ int64_t seq_length = GetJSONIntegerField(metadata, "seq_length",
+ "PagedAttentionKVCache checkpoint
import metadata");
+ TVM_FFI_ICHECK_GE(seq_length, 0)
+ << "PagedAttentionKVCache checkpoint import seq_length cannot be
negative.";
+ TVM_FFI_ICHECK_LE(seq_length, std::numeric_limits<int32_t>::max())
+ << "PagedAttentionKVCache checkpoint import seq_length exceeds int32
range.";
+ int64_t num_logical_pages = GetExpectedNumLogicalPages(seq_length);
+ int64_t source_num_total_pages = GetJSONIntegerField(
+ metadata, "num_total_pages", "PagedAttentionKVCache checkpoint import
metadata");
+ TVM_FFI_ICHECK_LE(num_logical_pages, source_num_total_pages)
+ << "PagedAttentionKVCache checkpoint metadata requires " <<
num_logical_pages
+ << " pages, but reports a source cache with only " <<
source_num_total_pages << " pages.";
+ TVM_FFI_ICHECK_LE(num_logical_pages, num_total_pages_)
+ << "PagedAttentionKVCache checkpoint import requires " <<
num_logical_pages
+ << " pages, but this cache only has " << num_total_pages_ << " pages.";
+ CheckCheckpointImportLogicalPages(metadata, seq_length);
+ CheckCheckpointImportGroups(metadata, num_logical_pages);
+ return seq_length;
+ }
+
+ void CheckPageGroupTensor(const Tensor& tensor, int64_t num_logical_pages,
+ const char* api_name) const {
+ std::string error_msg = std::string(api_name) +
+ " expects the tensor in layout "
+
"(1,num_logical_pages,2,num_kv_heads,page_size,qk_head_dim).";
+ TVM_FFI_ICHECK(tensor.defined()) << error_msg;
+ TVM_FFI_ICHECK(tensor.DataType() == kv_dtype_)
+ << error_msg << " The dtype mismatches, expected " << kv_dtype_ << ",
got "
+ << tensor.DataType() << ".";
+ TVM_FFI_ICHECK_EQ(tensor->ndim, 6) << error_msg;
+ TVM_FFI_ICHECK_EQ(tensor->shape[0], 1) << error_msg << " The group count
mismatches.";
+ TVM_FFI_ICHECK_EQ(tensor->shape[1], num_logical_pages)
+ << error_msg << " The number of logical pages mismatches.";
+ TVM_FFI_ICHECK_EQ(tensor->shape[2], 2) << error_msg << " The K/V axis
mismatches.";
+ TVM_FFI_ICHECK_EQ(tensor->shape[3], num_kv_heads_)
+ << error_msg << " The number of KV heads mismatches.";
+ TVM_FFI_ICHECK_EQ(tensor->shape[4], page_size_) << error_msg << " The page
size mismatches.";
+ TVM_FFI_ICHECK_EQ(tensor->shape[5], qk_head_dim_)
+ << error_msg << " The head dimension mismatches.";
+ }
+
+ void CheckExportPageGroupTensor(const Tensor& dst, int64_t
num_logical_pages) const {
+ CheckPageGroupTensor(dst, num_logical_pages, "ExportPageGroup");
+ }
+
+ void CheckImportPageGroupTensor(const Tensor& src, int64_t
num_logical_pages) const {
+ CheckPageGroupTensor(src, num_logical_pages, "ImportPageGroup");
+ }
+
+ ffi::json::Array MakeAttnKindsMetadata() const {
+ ffi::json::Array result;
+ for (int64_t layer_id = layer_id_begin_offset_; layer_id <
layer_id_end_offset_; ++layer_id) {
+ AttnKind attn_kind = attn_kinds_[layer_id];
+ result.push_back(ffi::String(AttnKindToString(attn_kind)));
+ }
+ return result;
+ }
+
+ ffi::json::Object MakeLayoutMetadata() const {
+ namespace json = tvm::ffi::json;
+ json::Object metadata;
+ metadata.Set("format_version", kPagedKVCacheCheckpointFormatVersion);
+ metadata.Set("cache_type", ffi::String(kPagedKVCacheCheckpointRuntime));
+ metadata.Set("page_size", page_size_);
+ metadata.Set("num_layers", num_layers_);
+ metadata.Set("layer_begin", layer_id_begin_offset_);
+ metadata.Set("layer_end", layer_id_end_offset_);
+ metadata.Set("num_qo_heads", num_qo_heads_);
+ metadata.Set("num_kv_heads", num_kv_heads_);
+ metadata.Set("qk_head_dim", qk_head_dim_);
+ metadata.Set("v_head_dim", v_head_dim_);
+ metadata.Set("reserved_num_seqs", reserved_num_seqs_);
+ metadata.Set("num_total_pages", num_total_pages_);
+ metadata.Set("prefill_chunk_size", prefill_chunk_size_);
+ metadata.Set("dtype", ffi::DLDataTypeToString(kv_dtype_));
+ metadata.Set("attn_kinds", MakeAttnKindsMetadata());
+ metadata.Set("rope_mode", ffi::String(RoPEModeToString(rope_mode_)));
+ metadata.Set("rotary_scale", rotary_scale_);
+ metadata.Set("rotary_theta", rotary_theta_);
+ metadata.Set("has_rope_ext_factors", rope_ext_factors_.has_value());
+ metadata.Set("support_sliding_window", support_sliding_window_);
+ metadata.Set("support_layer_sliding_window",
support_layer_sliding_window_);
+ metadata.Set("page_group_layout",
+
ffi::String("1,num_logical_pages,2,num_kv_heads,page_size,qk_head_dim"));
+ return metadata;
+ }
+
+ std::string GetLayoutDescriptor() const {
+ std::ostringstream os;
+ os.imbue(std::locale::classic());
+ os << std::setprecision(std::numeric_limits<double>::max_digits10);
+ os << "format_version=" << kPagedKVCacheCheckpointFormatVersion << ";";
+ os << "cache_type=" << kPagedKVCacheCheckpointRuntime << ";";
+ os << "page_size=" << page_size_ << ";";
+ os << "num_layers=" << num_layers_ << ";";
+ os << "layer_begin=" << layer_id_begin_offset_ << ";";
+ os << "layer_end=" << layer_id_end_offset_ << ";";
+ os << "num_qo_heads=" << num_qo_heads_ << ";";
+ os << "num_kv_heads=" << num_kv_heads_ << ";";
+ os << "qk_head_dim=" << qk_head_dim_ << ";";
+ os << "v_head_dim=" << v_head_dim_ << ";";
+ os << "dtype=" << std::string(ffi::DLDataTypeToString(kv_dtype_)) << ";";
+ os << "rope_mode=" << RoPEModeToString(rope_mode_) << ";";
+ os << "rotary_scale=" << rotary_scale_ << ";";
+ os << "rotary_theta=" << rotary_theta_ << ";";
+ os << "has_rope_ext_factors=" << rope_ext_factors_.has_value() << ";";
+ os << "support_sliding_window=" << support_sliding_window_ << ";";
+ os << "support_layer_sliding_window=" << support_layer_sliding_window_ <<
";";
+ os <<
"page_group_layout=1,num_logical_pages,2,num_kv_heads,page_size,qk_head_dim;";
+ os << "attn_kinds=";
+ for (int64_t layer_id = layer_id_begin_offset_; layer_id <
layer_id_end_offset_; ++layer_id) {
+ if (layer_id != layer_id_begin_offset_) {
+ os << ",";
+ }
+ os << AttnKindToString(attn_kinds_[layer_id]);
+ }
+ return os.str();
+ }
+
+ ffi::json::Array MakeLogicalPageMetadata(const Sequence& seq) const {
+ namespace json = tvm::ffi::json;
+ json::Array pages;
+ int64_t num_logical_pages = GetExpectedNumLogicalPages(seq.seq_length);
+ for (int64_t logical_page_index = 0; logical_page_index <
num_logical_pages;
+ ++logical_page_index) {
+ int64_t page_start = logical_page_index * page_size_;
+ int64_t page_length = std::min<int64_t>(page_size_, seq.seq_length -
page_start);
+ TVM_FFI_ICHECK_GT(page_length, 0);
+ json::Object page_json;
+ page_json.Set("logical_page_index", logical_page_index);
+ page_json.Set("start_pos", page_start);
+ page_json.Set("length", page_length);
+ pages.push_back(page_json);
+ }
+ return pages;
+ }
+
+ int64_t GetNumLogicalPages(const Sequence& seq) const {
+ int64_t num_pages = 0;
+ for (int32_t block_id : seq.GetBlockTrace(global_block_pool_)) {
+ num_pages += global_block_pool_[block_id].page_ids.size();
+ }
+ TVM_FFI_ICHECK_EQ(num_pages, GetExpectedNumLogicalPages(seq.seq_length))
+ << "PagedAttentionKVCache checkpointing found a page table that does
not match the "
+ "sequence length.";
+ return num_pages;
+ }
+
+ ffi::json::Array MakePageGroupMetadata(const Sequence& seq) const {
+ namespace json = tvm::ffi::json;
+ json::Array groups;
+ int64_t num_logical_pages = GetNumLogicalPages(seq);
+ for (int64_t local_layer = 0; local_layer < num_layers_; ++local_layer) {
+ json::Object group;
+ group.Set("group_index", local_layer);
+ group.Set("layer_begin", layer_id_begin_offset_ + local_layer);
+ group.Set("layer_end", layer_id_begin_offset_ + local_layer + 1);
+ group.Set("num_logical_pages", num_logical_pages);
+ group.Set("dtype", ffi::DLDataTypeToString(kv_dtype_));
+ group.Set("shape",
+ json::Array{1, num_logical_pages, 2, num_kv_heads_,
page_size_, qk_head_dim_});
+ group.Set("nbytes", GetCheckpointGroupNBytes(num_logical_pages));
+ groups.push_back(group);
+ }
+ return groups;
+ }
+
+ int64_t GetCheckpointBytesPerScalar() const {
+ return (static_cast<int64_t>(kv_dtype_.bits) * kv_dtype_.lanes + 7) / 8;
+ }
+
+ int64_t GetCheckpointBytesPerPage() const {
+ return 2 * num_kv_heads_ * page_size_ * qk_head_dim_ *
GetCheckpointBytesPerScalar();
+ }
+
+ int64_t GetCheckpointGroupNBytes(int64_t num_logical_pages) const {
+ return num_logical_pages * GetCheckpointBytesPerPage();
+ }
+
+ std::vector<int32_t> GetCheckpointPageIds(const Sequence& seq, const char*
operation) const {
+ std::vector<int32_t> page_ids;
+ page_ids.reserve(GetExpectedNumLogicalPages(seq.seq_length));
+ for (int32_t block_id : seq.GetBlockTrace(global_block_pool_)) {
+ const Block& block = global_block_pool_[block_id];
+ for (int32_t page_id : block.page_ids) {
+ TVM_FFI_ICHECK_GE(page_id, 0)
+ << "PagedAttentionKVCache " << operation << " found invalid page
id " << page_id << ".";
+ TVM_FFI_ICHECK_LT(page_id, num_total_pages_)
+ << "PagedAttentionKVCache " << operation << " found out-of-range
page id " << page_id
+ << ".";
+ page_ids.push_back(page_id);
+ }
+ }
+ TVM_FFI_ICHECK_EQ(static_cast<int64_t>(page_ids.size()),
+ GetExpectedNumLogicalPages(seq.seq_length))
+ << "PagedAttentionKVCache " << operation
+ << " found a page table that does not match the sequence length.";
+ return page_ids;
+ }
+
+ void CopyCheckpointPageRun(const Tensor& src, int64_t src_page_index, const
Tensor& dst,
+ int64_t dst_page_index, int64_t run_length) const
{
+ TVM_FFI_ICHECK_GT(run_length, 0);
+ int64_t bytes_per_page = GetCheckpointBytesPerPage();
+ Tensor src_pages =
+ src.CreateView({run_length, 2, num_kv_heads_, page_size_,
qk_head_dim_}, src->dtype,
+ static_cast<uint64_t>(src_page_index * bytes_per_page));
+ Tensor dst_pages =
+ dst.CreateView({run_length, 2, num_kv_heads_, page_size_,
qk_head_dim_}, dst->dtype,
+ static_cast<uint64_t>(dst_page_index * bytes_per_page));
+ DLTensor dst_pages_view = *dst_pages.operator->();
+ Tensor::CopyFromTo(src_pages.operator->(), &dst_pages_view,
compute_stream_);
+ }
+
+ void ExportPartialCheckpointPage(const Tensor& src, int64_t src_page_index,
const Tensor& dst,
+ int64_t dst_page_index, int64_t
valid_length) {
+ TVM_FFI_ICHECK_GT(valid_length, 0);
+ TVM_FFI_ICHECK_LT(valid_length, page_size_);
+
+ int64_t bytes_per_scalar = GetCheckpointBytesPerScalar();
+ int64_t bytes_per_page = GetCheckpointBytesPerPage();
+ if (!checkpoint_zero_page_host_.defined()) {
+ checkpoint_zero_page_host_ = Tensor::Empty({2, num_kv_heads_,
page_size_, qk_head_dim_},
+ kv_dtype_,
GetPreferredHostDevice(device_));
+ std::vector<uint8_t> zero_data(bytes_per_page, 0);
+ checkpoint_zero_page_host_.CopyFromBytes(zero_data.data(),
zero_data.size());
+ }
+
+ // Initialize the full destination page with one transfer, then overwrite
+ // only the valid prefixes. This makes every unused slot deterministic.
+ Tensor dst_page = dst.CreateView({2, num_kv_heads_, page_size_,
qk_head_dim_}, dst->dtype,
+ static_cast<uint64_t>(dst_page_index *
bytes_per_page));
+ DLTensor dst_page_view = *dst_page.operator->();
+ Tensor::CopyFromTo(checkpoint_zero_page_host_.operator->(),
&dst_page_view, compute_stream_);
+
+ for (int64_t kv_index = 0; kv_index < 2; ++kv_index) {
+ for (int64_t head_index = 0; head_index < num_kv_heads_; ++head_index) {
+ int64_t head_offset = (kv_index * num_kv_heads_ + head_index) *
page_size_ * qk_head_dim_;
+ uint64_t src_offset =
+ static_cast<uint64_t>(src_page_index * bytes_per_page +
head_offset * bytes_per_scalar);
+ uint64_t dst_offset =
+ static_cast<uint64_t>(dst_page_index * bytes_per_page +
head_offset * bytes_per_scalar);
+ Tensor src_valid = src.CreateView({valid_length, qk_head_dim_},
src->dtype, src_offset);
+ Tensor dst_valid = dst.CreateView({valid_length, qk_head_dim_},
dst->dtype, dst_offset);
+
+ DLTensor dst_valid_view = *dst_valid.operator->();
+ Tensor::CopyFromTo(src_valid.operator->(), &dst_valid_view,
compute_stream_);
+ }
+ }
+ }
+
/*! \brief Get a new free page and return its id. */
int32_t GetFreePage() {
// Find a page from the free page pools.
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
index ed86585e5d..6fc0b516bc 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
@@ -74,6 +74,12 @@ fattention_with_fuse_qkv = None
fattention_with_shared_kv = None
fis_empty = None
fdebug_get_kv = None
+fget_checkpoint_metadata = None
+fexport_page_group = None
+fprepare_import = None
+fimport_page_group = None
+ffinish_import = None
+fget_sequence_length = None
ftranspose_append = None
fcopy_cache = None
@@ -97,6 +103,8 @@ def set_global_func(head_dim, dtype,
layer_sliding_window_size=1024):
global fclear, fadd_sequence, fremove_sequence, ffork_sequence,
fenable_sliding_window_for_seq
global fpopn, fbegin_forward, fend_forward,
fcommit_accepted_token_tree_nodes
global fattention_with_fuse_qkv, fattention_with_shared_kv, fis_empty,
fdebug_get_kv
+ global fget_checkpoint_metadata, fexport_page_group
+ global fprepare_import, fimport_page_group, ffinish_import,
fget_sequence_length
global ftranspose_append, fcopy_cache, fattn_prefill, fattn_decode
global \
fattn_prefill_ragged, \
@@ -126,6 +134,14 @@ def set_global_func(head_dim, dtype,
layer_sliding_window_size=1024):
)
fis_empty = tvm.get_global_func("vm.builtin.attention_kv_cache_empty")
fdebug_get_kv =
tvm.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv")
+ fget_checkpoint_metadata = tvm.get_global_func(
+ "vm.builtin.attention_kv_cache_get_checkpoint_metadata"
+ )
+ fexport_page_group =
tvm.get_global_func("vm.builtin.attention_kv_cache_export_page_group")
+ fprepare_import =
tvm.get_global_func("vm.builtin.attention_kv_cache_prepare_import")
+ fimport_page_group =
tvm.get_global_func("vm.builtin.attention_kv_cache_import_page_group")
+ ffinish_import =
tvm.get_global_func("vm.builtin.attention_kv_cache_finish_import")
+ fget_sequence_length =
tvm.get_global_func("vm.builtin.attention_kv_cache_get_sequence_length")
target = tvm.target.Target.from_device(device)
cache_key = (
@@ -302,6 +318,66 @@ def verify_cached_kv(kv_cache, seq_ids, expected_k,
expected_v):
tvm.testing.assert_allclose(values.numpy(), values_expected,
rtol=1e-3, atol=1e-3)
+def verify_exported_page_groups(kv_cache, seq_id):
+ metadata = json.loads(fget_checkpoint_metadata(kv_cache, seq_id))
+ seq_length = metadata["seq_length"]
+ keys = tvm.runtime.empty(
+ (num_layers, seq_length, num_kv_heads, head_dim), dtype=dtype,
device=device
+ )
+ values = tvm.runtime.empty(
+ (num_layers, seq_length, num_kv_heads, head_dim), dtype=dtype,
device=device
+ )
+ fdebug_get_kv(kv_cache, seq_id, 0, seq_length, keys, values)
+ keys_np = keys.numpy()
+ values_np = values.numpy()
+
+ for group in metadata["groups"]:
+ group_data = tvm.runtime.empty(tuple(group["shape"]), dtype=dtype,
device=device)
+ fexport_page_group(kv_cache, seq_id, group["group_index"], group_data)
+ group_np = group_data.numpy()
+ layer = group["group_index"]
+ for page in metadata["logical_pages"]:
+ logical_page_index = page["logical_page_index"]
+ start_pos = page["start_pos"]
+ length = page["length"]
+ exported_k = group_np[0, logical_page_index, 0, :, :length,
:].transpose(1, 0, 2)
+ exported_v = group_np[0, logical_page_index, 1, :, :length,
:].transpose(1, 0, 2)
+ tvm.testing.assert_allclose(
+ exported_k, keys_np[layer, start_pos : start_pos + length],
rtol=1e-3, atol=1e-3
+ )
+ tvm.testing.assert_allclose(
+ exported_v, values_np[layer, start_pos : start_pos + length],
rtol=1e-3, atol=1e-3
+ )
+
+
+def export_page_groups(kv_cache, metadata):
+ groups = []
+ for group in metadata["groups"]:
+ group_data = tvm.runtime.empty(tuple(group["shape"]), dtype=dtype,
device=device)
+ fexport_page_group(kv_cache, metadata["seq_id"], group["group_index"],
group_data)
+ groups.append(group_data)
+ return groups
+
+
+def verify_debug_kv_equal(lhs_cache, rhs_cache, seq_id, seq_length):
+ lhs_keys = tvm.runtime.empty(
+ (num_layers, seq_length, num_kv_heads, head_dim), dtype=dtype,
device=device
+ )
+ lhs_values = tvm.runtime.empty(
+ (num_layers, seq_length, num_kv_heads, head_dim), dtype=dtype,
device=device
+ )
+ rhs_keys = tvm.runtime.empty(
+ (num_layers, seq_length, num_kv_heads, head_dim), dtype=dtype,
device=device
+ )
+ rhs_values = tvm.runtime.empty(
+ (num_layers, seq_length, num_kv_heads, head_dim), dtype=dtype,
device=device
+ )
+ fdebug_get_kv(lhs_cache, seq_id, 0, seq_length, lhs_keys, lhs_values)
+ fdebug_get_kv(rhs_cache, seq_id, 0, seq_length, rhs_keys, rhs_values)
+ tvm.testing.assert_allclose(lhs_keys.numpy(), rhs_keys.numpy(), rtol=1e-3,
atol=1e-3)
+ tvm.testing.assert_allclose(lhs_values.numpy(), rhs_values.numpy(),
rtol=1e-3, atol=1e-3)
+
+
def f_apply_rotary(x, offset, scale, theta, offset_list: list[int] | None =
None):
# x: (N, H, D)
assert len(x.shape) == 3
@@ -726,6 +802,70 @@ def
test_paged_attention_kv_cache_prefill_and_decode(kv_cache_and_config):
apply_attention(kv_cache, rope_mode, batch, cached_k, cached_v)
+def test_paged_attention_kv_cache_export_page_group():
+ global head_dim, sm_scale, dtype
+ head_dim = 64
+ dtype = "float32"
+ sm_scale = head_dim ** (-0.5)
+ set_global_func(head_dim, dtype)
+ kv_cache = create_kv_cache(head_dim, dtype, RopeMode.NONE, False)
+
+ cached_k = {}
+ cached_v = {}
+ apply_attention(kv_cache, RopeMode.NONE, [(0, page_size * 2 + 3)],
cached_k, cached_v)
+ verify_exported_page_groups(kv_cache, 0)
+
+ # Reuse a page whose unused slots contain old KV values. A checkpoint must
+ # zero the unused tail rather than expose those stale values.
+ fclear(kv_cache)
+ cached_k = {}
+ cached_v = {}
+ apply_attention(kv_cache, RopeMode.NONE, [(0, page_size)], cached_k,
cached_v)
+ fclear(kv_cache)
+ cached_k = {}
+ cached_v = {}
+ apply_attention(kv_cache, RopeMode.NONE, [(0, 1)], cached_k, cached_v)
+ metadata = json.loads(fget_checkpoint_metadata(kv_cache, 0))
+ group_data = export_page_groups(kv_cache, metadata)[0].numpy()
+ np.testing.assert_array_equal(group_data[0, 0, :, :, 1:, :], 0)
+
+
+def test_paged_attention_kv_cache_import_page_group_round_trip():
+ global head_dim, sm_scale, dtype
+ head_dim = 64
+ dtype = "float32"
+ sm_scale = head_dim ** (-0.5)
+ set_global_func(head_dim, dtype)
+ src_cache = create_kv_cache(head_dim, dtype, RopeMode.NONE, False)
+
+ cached_k = {}
+ cached_v = {}
+ apply_attention(src_cache, RopeMode.NONE, [(0, page_size * 2 + 3)],
cached_k, cached_v)
+
+ metadata_json = fget_checkpoint_metadata(src_cache, 0)
+ metadata = json.loads(metadata_json)
+ groups = export_page_groups(src_cache, metadata)
+
+ dst_cache = create_kv_cache(head_dim, dtype, RopeMode.NONE, False)
+ fprepare_import(dst_cache, 0, metadata_json)
+ for group, group_data in zip(metadata["groups"], groups):
+ fimport_page_group(dst_cache, 0, group["group_index"], group_data)
+ ffinish_import(dst_cache, 0)
+ assert fget_sequence_length(dst_cache, 0) == metadata["seq_length"]
+
+ verify_debug_kv_equal(src_cache, dst_cache, 0, metadata["seq_length"])
+
+ # Continue decoding from both caches with identical inputs. This exercises
+ # the restored page table and the partially filled final page.
+ dst_cached_k = {0: cached_k[0].copy()}
+ dst_cached_v = {0: cached_v[0].copy()}
+ random_state = np.random.get_state()
+ apply_attention(src_cache, RopeMode.NONE, [(0, 1)], cached_k, cached_v)
+ np.random.set_state(random_state)
+ apply_attention(dst_cache, RopeMode.NONE, [(0, 1)], dst_cached_k,
dst_cached_v)
+ verify_debug_kv_equal(src_cache, dst_cache, 0, metadata["seq_length"] + 1)
+
+
def test_paged_attention_kv_cache_remove_sequence(kv_cache_and_config):
kv_cache, rope_mode, support_sliding_window = kv_cache_and_config
if support_sliding_window and rope_mode == RopeMode.NORMAL:
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_metadata.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_metadata.py
new file mode 100644
index 0000000000..25e642d795
--- /dev/null
+++
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_metadata.py
@@ -0,0 +1,417 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import json
+
+import pytest
+import tvm_ffi
+from tvm_ffi import Shape
+
+import tvm
+import tvm.testing
+from tvm.error import InternalError
+from tvm.relax.frontend.nn.llm.kv_cache import AttnKind, RopeMode
+
+reserved_nseq = 4
+maximum_total_seq_length = 128
+prefill_chunk_size = 64
+page_size = 16
+num_layers = 4
+num_qo_heads = 32
+num_kv_heads = 4
+head_dim = 64
+rope_scale = 1.0
+rope_theta = 1e4
+device = tvm.cpu()
+
+
+def _nop(*args):
+ return None
+
+
+def create_kv_cache(
+ *,
+ dtype="float16",
+ head_dim_value=head_dim,
+ v_head_dim_value=None,
+ page_size_value=page_size,
+ num_layers_value=num_layers,
+ rope_mode=RopeMode.NORMAL,
+ attn_kind=AttnKind.MHA,
+ support_sliding_window=False,
+ reserved_nseq_value=reserved_nseq,
+ maximum_total_seq_length_value=maximum_total_seq_length,
+ prefill_chunk_size_value=prefill_chunk_size,
+ rope_ext_factors=None,
+):
+ fcreate = tvm.get_global_func("vm.builtin.paged_attention_kv_cache_create")
+ dummy_func = tvm.runtime.convert(_nop)
+ return fcreate(
+ tvm_ffi.Shape(
+ [
+ reserved_nseq_value,
+ maximum_total_seq_length_value,
+ prefill_chunk_size_value,
+ page_size_value,
+ int(support_sliding_window),
+ ]
+ ),
+ tvm_ffi.Shape([0, num_layers_value]),
+ num_qo_heads,
+ num_kv_heads,
+ head_dim_value,
+ head_dim_value if v_head_dim_value is None else v_head_dim_value,
+ tvm_ffi.Shape([int(attn_kind) for _ in range(num_layers_value)]),
+ False, # enable_kv_transfer
+ int(rope_mode),
+ rope_scale,
+ rope_theta,
+ rope_ext_factors,
+ tvm.runtime.empty((), dtype, device=device),
+ dummy_func, # f_transpose_append_mha
+ None, # f_transpose_append_mla
+ [], # f_attention_prefill_ragged
+ [], # f_attention_prefill
+ [], # f_attention_decode
+ [], # f_attention_prefill_sliding_window
+ [], # f_attention_decode_sliding_window
+ [], # f_attention_prefill_with_tree_mask_paged_kv
+ [], # f_attention_prefill_with_tree_mask
+ [], # f_mla_prefill
+ [dummy_func], # f_merge_inplace
+ dummy_func, # f_split_rotary
+ dummy_func, # f_copy_single_page
+ dummy_func, # f_debug_get_kv
+ dummy_func, # f_compact_copy
+ )
+
+
+def append_tokens(kv_cache, seq_id=0, append_length=page_size + 1):
+ fadd_sequence = tvm.get_global_func("vm.builtin.kv_state_add_sequence")
+ fbegin_forward = tvm.get_global_func("vm.builtin.kv_state_begin_forward")
+ fend_forward = tvm.get_global_func("vm.builtin.kv_state_end_forward")
+ fadd_sequence(kv_cache, seq_id)
+ fbegin_forward(kv_cache, Shape([seq_id]), Shape([append_length]), None)
+ fend_forward(kv_cache)
+
+
+def test_checkpoint_metadata_reports_layout_pages_and_groups():
+ fget_checkpoint_metadata = tvm.get_global_func(
+ "vm.builtin.attention_kv_cache_get_checkpoint_metadata"
+ )
+ fget_layout_hash =
tvm.get_global_func("vm.builtin.attention_kv_cache_get_layout_hash")
+ fexport_page_group =
tvm.get_global_func("vm.builtin.attention_kv_cache_export_page_group")
+ fprepare_import =
tvm.get_global_func("vm.builtin.attention_kv_cache_prepare_import")
+ fimport_page_group =
tvm.get_global_func("vm.builtin.attention_kv_cache_import_page_group")
+ ffinish_import =
tvm.get_global_func("vm.builtin.attention_kv_cache_finish_import")
+ fget_sequence_length =
tvm.get_global_func("vm.builtin.attention_kv_cache_get_sequence_length")
+
+ kv_cache = create_kv_cache()
+ append_tokens(kv_cache)
+ metadata_json = fget_checkpoint_metadata(kv_cache, 0)
+ metadata = json.loads(metadata_json)
+
+ assert metadata["format_version"] == 1
+ assert metadata["cache_type"] == "relax.vm.PagedAttentionKVCache"
+ assert metadata["layout_hash"] == fget_layout_hash(kv_cache)
+ assert metadata["seq_id"] == 0
+ assert metadata["seq_length"] == page_size + 1
+ assert metadata["page_size"] == page_size
+ assert metadata["dtype"] == "float16"
+ assert metadata["layer_begin"] == 0
+ assert metadata["layer_end"] == num_layers
+ assert metadata["num_kv_heads"] == num_kv_heads
+ assert metadata["qk_head_dim"] == head_dim
+ assert metadata["v_head_dim"] == head_dim
+ assert metadata["reserved_num_seqs"] == reserved_nseq
+ assert metadata["attn_kinds"] == ["mha"] * num_layers
+ assert (
+ metadata["page_group_layout"] ==
"1,num_logical_pages,2,num_kv_heads,page_size,qk_head_dim"
+ )
+ assert "blocks" not in metadata
+ assert len(metadata["logical_pages"]) == 2
+ assert metadata["logical_pages"][0]["start_pos"] == 0
+ assert metadata["logical_pages"][0]["length"] == page_size
+ assert metadata["logical_pages"][1]["start_pos"] == page_size
+ assert metadata["logical_pages"][1]["length"] == 1
+ assert "page_id" not in metadata["logical_pages"][0]
+ assert "block_index" not in metadata["logical_pages"][0]
+ assert len(metadata["groups"]) == num_layers
+ assert metadata["groups"][0]["layer_begin"] == 0
+ assert metadata["groups"][0]["layer_end"] == 1
+ assert metadata["groups"][0]["num_logical_pages"] == 2
+ assert metadata["groups"][0]["dtype"] == "float16"
+ assert metadata["groups"][0]["shape"] == [1, 2, 2, num_kv_heads,
page_size, head_dim]
+
+ exported_groups = []
+ for group_metadata in metadata["groups"]:
+ group = tvm.runtime.empty(tuple(group_metadata["shape"]), "float16",
device=device)
+ fexport_page_group(kv_cache, 0, group_metadata["group_index"], group)
+ exported_groups.append(group)
+
+ import_cache = create_kv_cache(
+ reserved_nseq_value=reserved_nseq * 2,
+ maximum_total_seq_length_value=maximum_total_seq_length * 2,
+ prefill_chunk_size_value=prefill_chunk_size // 2,
+ )
+ fprepare_import(import_cache, 0, metadata_json)
+ with pytest.raises(InternalError, match="until checkpoint import is
finished"):
+ fget_sequence_length(import_cache, 0)
+ for group_metadata, group in zip(metadata["groups"], exported_groups):
+ fimport_page_group(import_cache, 0, group_metadata["group_index"],
group)
+ ffinish_import(import_cache, 0)
+ assert fget_sequence_length(import_cache, 0) == page_size + 1
+
+
+def test_checkpoint_layout_hash_is_stable_and_layout_sensitive():
+ fget_layout_hash =
tvm.get_global_func("vm.builtin.attention_kv_cache_get_layout_hash")
+
+ kv_cache = create_kv_cache()
+ same_layout = create_kv_cache()
+ different_page_size = create_kv_cache(page_size_value=page_size * 2)
+ different_num_layers = create_kv_cache(num_layers_value=2)
+ different_head_dim = create_kv_cache(head_dim_value=128)
+ different_dtype = create_kv_cache(dtype="float32")
+ different_rope = create_kv_cache(rope_mode=RopeMode.NONE)
+ different_operational_limits = create_kv_cache(
+ reserved_nseq_value=reserved_nseq * 2,
+ maximum_total_seq_length_value=maximum_total_seq_length * 2,
+ prefill_chunk_size_value=prefill_chunk_size // 2,
+ )
+
+ layout_hash = fget_layout_hash(kv_cache)
+ assert layout_hash == fget_layout_hash(kv_cache)
+ assert layout_hash == fget_layout_hash(same_layout)
+ assert layout_hash != fget_layout_hash(different_page_size)
+ assert layout_hash != fget_layout_hash(different_num_layers)
+ assert layout_hash != fget_layout_hash(different_head_dim)
+ assert layout_hash != fget_layout_hash(different_dtype)
+ assert layout_hash != fget_layout_hash(different_rope)
+ assert layout_hash == fget_layout_hash(different_operational_limits)
+
+
+def test_checkpoint_metadata_rejects_unsupported_layouts_and_sequence_ids():
+ fget_checkpoint_metadata = tvm.get_global_func(
+ "vm.builtin.attention_kv_cache_get_checkpoint_metadata"
+ )
+ fget_layout_hash =
tvm.get_global_func("vm.builtin.attention_kv_cache_get_layout_hash")
+ fexport_page_group =
tvm.get_global_func("vm.builtin.attention_kv_cache_export_page_group")
+
+ sliding_cache = create_kv_cache(support_sliding_window=True)
+ append_tokens(sliding_cache)
+ mla_cache = create_kv_cache(attn_kind=AttnKind.MLA)
+ asymmetric_cache = create_kv_cache(v_head_dim_value=head_dim // 2)
+ rope_ext_cache = create_kv_cache(
+ rope_ext_factors=tvm.runtime.empty((head_dim // 2,), "float32",
device=device)
+ )
+ dst = tvm.runtime.empty((1, 1, 2, num_kv_heads, page_size, head_dim),
"float16", device=device)
+
+ with pytest.raises(InternalError, match="sliding-window"):
+ fget_layout_hash(sliding_cache)
+ with pytest.raises(InternalError, match="sliding-window"):
+ fget_checkpoint_metadata(sliding_cache, 0)
+ with pytest.raises(InternalError, match="sliding-window"):
+ fexport_page_group(sliding_cache, 0, 0, dst)
+ with pytest.raises(InternalError, match="sequence id 0"):
+ fget_checkpoint_metadata(create_kv_cache(), 1)
+ with pytest.raises(InternalError, match="full-context MHA/GQA"):
+ fget_layout_hash(mla_cache)
+ with pytest.raises(InternalError, match="full-context MHA/GQA"):
+ fexport_page_group(mla_cache, 0, 0, dst)
+ with pytest.raises(InternalError, match="qk_head_dim to equal v_head_dim"):
+ fget_layout_hash(asymmetric_cache)
+ with pytest.raises(InternalError, match="RoPE extension factors"):
+ fget_layout_hash(rope_ext_cache)
+
+ tree_cache = create_kv_cache()
+ tvm.get_global_func("vm.builtin.kv_state_add_sequence")(tree_cache, 0)
+ tvm.get_global_func("vm.builtin.kv_state_begin_forward")(
+ tree_cache, Shape([0]), Shape([2]), Shape([-1, 0])
+ )
+ with pytest.raises(InternalError, match="committed token-chain state"):
+ fexport_page_group(tree_cache, 0, 0, dst)
+
+
+def test_checkpoint_export_page_group_validates_group_shape():
+ fget_checkpoint_metadata = tvm.get_global_func(
+ "vm.builtin.attention_kv_cache_get_checkpoint_metadata"
+ )
+ fexport_page_group =
tvm.get_global_func("vm.builtin.attention_kv_cache_export_page_group")
+
+ kv_cache = create_kv_cache()
+ append_tokens(kv_cache)
+ metadata = json.loads(fget_checkpoint_metadata(kv_cache, 0))
+ shape = metadata["groups"][0]["shape"]
+
+ with pytest.raises(InternalError, match="group id"):
+ fexport_page_group(
+ kv_cache,
+ 0,
+ num_layers,
+ tvm.runtime.empty(tuple(shape), "float16", device=device),
+ )
+
+ bad_shape = shape.copy()
+ bad_shape[-1] += 1
+ with pytest.raises(InternalError, match="ExportPageGroup expects"):
+ fexport_page_group(
+ kv_cache,
+ 0,
+ 0,
+ tvm.runtime.empty(tuple(bad_shape), "float16", device=device),
+ )
+
+ with pytest.raises(InternalError, match="dtype mismatches"):
+ fexport_page_group(
+ kv_cache,
+ 0,
+ 0,
+ tvm.runtime.empty(tuple(shape), "float32", device=device),
+ )
+
+
+def test_checkpoint_prepare_import_validates_metadata():
+ fget_checkpoint_metadata = tvm.get_global_func(
+ "vm.builtin.attention_kv_cache_get_checkpoint_metadata"
+ )
+ fprepare_import =
tvm.get_global_func("vm.builtin.attention_kv_cache_prepare_import")
+
+ kv_cache = create_kv_cache()
+ append_tokens(kv_cache)
+ metadata_json = fget_checkpoint_metadata(kv_cache, 0)
+ metadata = json.loads(metadata_json)
+
+ with pytest.raises(InternalError, match="sequence id 0"):
+ fprepare_import(create_kv_cache(), 1, metadata_json)
+
+ with pytest.raises(InternalError, match="dtype"):
+ fprepare_import(create_kv_cache(dtype="float32"), 0, metadata_json)
+
+ bad_length = json.loads(metadata_json)
+ bad_length["seq_length"] = page_size * 2 + 1
+ with pytest.raises(InternalError, match="sequence length"):
+ fprepare_import(create_kv_cache(), 0, json.dumps(bad_length))
+
+ bad_group = json.loads(metadata_json)
+ bad_group["groups"][0]["shape"][-1] += 1
+ with pytest.raises(InternalError, match="shape"):
+ fprepare_import(create_kv_cache(), 0, json.dumps(bad_group))
+
+ bad_nbytes = json.loads(metadata_json)
+ bad_nbytes["groups"][0]["nbytes"] += 1
+ with pytest.raises(InternalError, match="nbytes"):
+ fprepare_import(create_kv_cache(), 0, json.dumps(bad_nbytes))
+
+ bad_version = json.loads(metadata_json)
+ bad_version["format_version"] += 1
+ with pytest.raises(InternalError, match="format_version"):
+ fprepare_import(create_kv_cache(), 0, json.dumps(bad_version))
+
+ with pytest.raises(InternalError, match="only has 1 pages"):
+ fprepare_import(
+ create_kv_cache(maximum_total_seq_length_value=0),
+ 0,
+ metadata_json,
+ )
+
+ metadata["layout_hash"] = "bad-layout-hash"
+ with pytest.raises(InternalError, match="layout hash mismatch"):
+ fprepare_import(create_kv_cache(), 0, json.dumps(metadata))
+
+
+def test_checkpoint_import_page_group_validates_group_shape():
+ fget_checkpoint_metadata = tvm.get_global_func(
+ "vm.builtin.attention_kv_cache_get_checkpoint_metadata"
+ )
+ fprepare_import =
tvm.get_global_func("vm.builtin.attention_kv_cache_prepare_import")
+ fimport_page_group =
tvm.get_global_func("vm.builtin.attention_kv_cache_import_page_group")
+
+ kv_cache = create_kv_cache()
+ append_tokens(kv_cache)
+ metadata_json = fget_checkpoint_metadata(kv_cache, 0)
+ metadata = json.loads(metadata_json)
+ shape = metadata["groups"][0]["shape"]
+
+ import_cache = create_kv_cache()
+ fprepare_import(import_cache, 0, metadata_json)
+
+ with pytest.raises(InternalError, match="group id"):
+ fimport_page_group(
+ import_cache,
+ 0,
+ num_layers,
+ tvm.runtime.empty(tuple(shape), "float16", device=device),
+ )
+
+ bad_shape = shape.copy()
+ bad_shape[-1] += 1
+ with pytest.raises(InternalError, match="ImportPageGroup expects"):
+ fimport_page_group(
+ import_cache,
+ 0,
+ 0,
+ tvm.runtime.empty(tuple(bad_shape), "float16", device=device),
+ )
+
+ with pytest.raises(InternalError, match="dtype mismatches"):
+ fimport_page_group(
+ import_cache,
+ 0,
+ 0,
+ tvm.runtime.empty(tuple(shape), "float32", device=device),
+ )
+
+
+def test_checkpoint_import_requires_all_groups_and_explicit_finish():
+ fget_checkpoint_metadata = tvm.get_global_func(
+ "vm.builtin.attention_kv_cache_get_checkpoint_metadata"
+ )
+ fexport_page_group =
tvm.get_global_func("vm.builtin.attention_kv_cache_export_page_group")
+ fprepare_import =
tvm.get_global_func("vm.builtin.attention_kv_cache_prepare_import")
+ fimport_page_group =
tvm.get_global_func("vm.builtin.attention_kv_cache_import_page_group")
+ ffinish_import =
tvm.get_global_func("vm.builtin.attention_kv_cache_finish_import")
+ fbegin_forward = tvm.get_global_func("vm.builtin.kv_state_begin_forward")
+
+ source_cache = create_kv_cache()
+ append_tokens(source_cache)
+ metadata_json = fget_checkpoint_metadata(source_cache, 0)
+ metadata = json.loads(metadata_json)
+ groups = []
+ for group_metadata in metadata["groups"]:
+ group = tvm.runtime.empty(tuple(group_metadata["shape"]), "float16",
device=device)
+ fexport_page_group(source_cache, 0, group_metadata["group_index"],
group)
+ groups.append(group)
+
+ import_cache = create_kv_cache()
+ fprepare_import(import_cache, 0, metadata_json)
+ fimport_page_group(import_cache, 0, 0, groups[0])
+
+ with pytest.raises(InternalError, match="already imported"):
+ fimport_page_group(import_cache, 0, 0, groups[0])
+ with pytest.raises(InternalError, match="missing group 1"):
+ ffinish_import(import_cache, 0)
+ with pytest.raises(InternalError, match="before checkpoint import is
finished"):
+ fbegin_forward(import_cache, Shape([0]), Shape([1]), None)
+
+ for group_id in range(1, num_layers):
+ fimport_page_group(import_cache, 0, group_id, groups[group_id])
+ ffinish_import(import_cache, 0)
+ with pytest.raises(InternalError, match="has not been prepared"):
+ ffinish_import(import_cache, 0)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
