gemini-code-assist[bot] commented on code in PR #20035:
URL: https://github.com/apache/tvm/pull/20035#discussion_r3613465060
##########
src/runtime/vm/paged_kv_cache.cc:
##########
@@ -1721,6 +1919,386 @@ 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(!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.";
Review Comment:
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()`.
```c
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.";
```
##########
src/runtime/vm/paged_kv_cache.cc:
##########
@@ -31,7 +32,10 @@
#include <tvm/support/cuda/nvtx.h>
#include <algorithm>
+#include <iomanip>
+#include <limits>
#include <numeric>
+#include <sstream>
Review Comment:
Include `<locale>` to support `std::locale::classic()` when formatting the
layout descriptor stream.
```suggestion
#include <iomanip>
#include <limits>
#include <locale>
#include <numeric>
#include <sstream>
```
##########
src/runtime/vm/paged_kv_cache.cc:
##########
@@ -1721,6 +1919,386 @@ 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(!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.";
+ }
+
+ void CheckCheckpointSequenceSupported(int64_t seq_id) const {
+ CheckCheckpointLayoutSupported();
+ 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>();
+ }
+
+ 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 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";
+ CheckJSONStringField(metadata, "cacheType",
kPagedKVCacheCheckpointRuntime, context);
+ CheckJSONIntegerField(metadata, "pageSize", page_size_, context);
+ CheckJSONIntegerField(metadata, "numLayers", num_layers_, context);
+ CheckJSONIntegerField(metadata, "layerBegin", layer_id_begin_offset_,
context);
+ CheckJSONIntegerField(metadata, "layerEnd", layer_id_end_offset_, context);
+ CheckJSONIntegerField(metadata, "numQOHeads", num_qo_heads_, context);
+ CheckJSONIntegerField(metadata, "numKVHeads", num_kv_heads_, context);
+ CheckJSONIntegerField(metadata, "qkHeadDim", qk_head_dim_, context);
+ CheckJSONIntegerField(metadata, "vHeadDim", v_head_dim_, context);
+ CheckJSONIntegerField(metadata, "numTotalPages", num_total_pages_,
context);
+ CheckJSONIntegerField(metadata, "prefillChunkSize", prefill_chunk_size_,
context);
+ CheckJSONStringField(metadata, "dtype",
std::string(ffi::DLDataTypeToString(kv_dtype_)),
+ context);
+ CheckJSONStringField(metadata, "ropeMode", RoPEModeToString(rope_mode_),
context);
+ CheckJSONBoolField(metadata, "hasRopeExtFactors",
rope_ext_factors_.has_value(), context);
+ CheckJSONBoolField(metadata, "supportSlidingWindow",
support_sliding_window_, context);
+ CheckJSONBoolField(metadata, "supportLayerSlidingWindow",
support_layer_sliding_window_,
+ context);
+ CheckJSONStringField(metadata, "pageTensorLayout",
+
"num_total_pages,2,num_kv_heads,page_size,qk_head_dim", context);
+
+ ffi::String expected_layout_hash = GetLayoutHash();
+ std::string layout_hash = GetJSONStringField(metadata, "layoutHash",
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, "attnKinds",
context);
+ TVM_FFI_ICHECK_EQ(attn_kinds.size(), num_layers_)
+ << context << " field \"attnKinds\" 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 \"attnKinds\" 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,
"logicalPages", context);
+ TVM_FFI_ICHECK_EQ(static_cast<int64_t>(logical_pages.size()),
num_logical_pages)
+ << "PagedAttentionKVCache checkpoint import sequence length mismatch:
seqLength "
+ << 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, "logicalPageIndex", i, context);
+ CheckJSONIntegerField(page, "startPos", 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
seqLength "
+ << 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, "groupIndex", local_layer, context);
+ CheckJSONIntegerField(group, "layerBegin", layer_id_begin_offset_ +
local_layer, context);
+ CheckJSONIntegerField(group, "layerEnd", layer_id_begin_offset_ +
local_layer + 1, context);
+ CheckJSONIntegerField(group, "numLogicalPages", num_logical_pages,
context);
+ CheckJSONStringField(group, "dtype",
std::string(ffi::DLDataTypeToString(kv_dtype_)),
+ context);
+ CheckPageGroupMetadataShape(group, num_logical_pages, context);
+ }
+ }
+
+ int64_t CheckCheckpointImportMetadata(int64_t seq_id, const
ffi::json::Object& metadata) const {
+ CheckCheckpointImportLayout(metadata);
+ CheckJSONIntegerField(metadata, "seqId", seq_id,
+ "PagedAttentionKVCache checkpoint import metadata");
+ int64_t seq_length = GetJSONIntegerField(metadata, "seqLength",
+ "PagedAttentionKVCache checkpoint
import metadata");
+ TVM_FFI_ICHECK_GE(seq_length, 0)
+ << "PagedAttentionKVCache checkpoint import seqLength cannot be
negative.";
+ TVM_FFI_ICHECK_LE(seq_length, std::numeric_limits<int32_t>::max())
+ << "PagedAttentionKVCache checkpoint import seqLength exceeds int32
range.";
+ int64_t num_logical_pages = GetExpectedNumLogicalPages(seq_length);
+ 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("cacheType", ffi::String(kPagedKVCacheCheckpointRuntime));
+ metadata.Set("pageSize", page_size_);
+ metadata.Set("numLayers", num_layers_);
+ metadata.Set("layerBegin", layer_id_begin_offset_);
+ metadata.Set("layerEnd", layer_id_end_offset_);
+ metadata.Set("numQOHeads", num_qo_heads_);
+ metadata.Set("numKVHeads", num_kv_heads_);
+ metadata.Set("qkHeadDim", qk_head_dim_);
+ metadata.Set("vHeadDim", v_head_dim_);
+ metadata.Set("numTotalPages", num_total_pages_);
+ metadata.Set("prefillChunkSize", prefill_chunk_size_);
+ metadata.Set("dtype", ffi::DLDataTypeToString(kv_dtype_));
+ metadata.Set("attnKinds", MakeAttnKindsMetadata());
+ metadata.Set("ropeMode", ffi::String(RoPEModeToString(rope_mode_)));
+ metadata.Set("rotaryScale", rotary_scale_);
+ metadata.Set("rotaryTheta", rotary_theta_);
+ metadata.Set("hasRopeExtFactors", rope_ext_factors_.has_value());
+ metadata.Set("supportSlidingWindow", support_sliding_window_);
+ metadata.Set("supportLayerSlidingWindow", support_layer_sliding_window_);
+ metadata.Set("pageTensorLayout",
+
ffi::String("num_total_pages,2,num_kv_heads,page_size,qk_head_dim"));
+ if (!pages_.empty()) {
+ metadata.Set("pageTensorShape", ShapeToJSON(pages_[0]->shape,
pages_[0]->ndim));
+ }
+ return metadata;
+ }
+
+ 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_ << ";";
Review Comment:
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.
```suggestion
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_ << ";";
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
