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 e345029ced [Fix][Runtime] Keep sliding-window page metadata separate
from global attention (#20699)
e345029ced is described below
commit e345029cedf0be76a087fcad9c89a97bed4f7204
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Fri Oct 23 20:49:03 2026 -0700
[Fix][Runtime] Keep sliding-window page metadata separate from global
attention (#20699)
`PlainPagedKVCacheAuxDataManager` reuses the same device arrays for
global and per-layer sliding attention. Preparing the sliding page table
therefore overwrites the global one before either layer uses it. In
Gemma 4, a global layer then attends to only the sliding window and
loses the earlier context.
Give the plain manager separate storage for the sliding page indptr,
page indices and length info, and select it at the per-layer sliding
copy calls. Allocate it only when the cache has sliding layers. The
cached manager already assigns each copy its own region.
---
src/runtime/vm/attn_utils.h | 59 +++++++++++++++++-------
src/runtime/vm/paged_kv_cache.cc | 9 ++--
tests/cpp/runtime/attn_aux_data_test.cc | 79 +++++++++++++++++++++++++++++++++
3 files changed, 127 insertions(+), 20 deletions(-)
diff --git a/src/runtime/vm/attn_utils.h b/src/runtime/vm/attn_utils.h
index 77ee094ca7..4a9f28c560 100644
--- a/src/runtime/vm/attn_utils.h
+++ b/src/runtime/vm/attn_utils.h
@@ -497,9 +497,11 @@ class PagedKVCacheAuxDataManager {
/*! \brief Copy the indptr array of append lengths after coalescing. (see
GetChunkedBlockIds) */
virtual ffi::Tensor CopyQOIndptrOnDepthAsync(HostMemoryVector* data, int
depth) = 0;
/*! \brief Copy the indptr array of page table. */
- virtual ffi::Tensor CopyPageIndptrOnDepthAsync(HostMemoryVector* data, int
depth) = 0;
+ virtual ffi::Tensor CopyPageIndptrOnDepthAsync(HostMemoryVector* data, int
depth,
+ bool layer_sliding_window =
false) = 0;
/*! \brief Copy the indices array of page table. */
- virtual ffi::Tensor CopyPageIndicesOnDepthAsync(HostMemoryVector* data, int
depth) = 0;
+ virtual ffi::Tensor CopyPageIndicesOnDepthAsync(HostMemoryVector* data, int
depth,
+ bool layer_sliding_window =
false) = 0;
/*! \brief Copy the array of KV slot number used in the last page of the
seq. */
virtual ffi::Tensor CopyLastPageLenOnDepthAsync(HostMemoryVector* data, int
depth) = 0;
/*!
@@ -514,7 +516,8 @@ class PagedKVCacheAuxDataManager {
*/
virtual ffi::Tensor CopyLengthInfoOnDepthAsync(HostMemoryVector*
last_page_len,
HostMemoryVector*
sliding_window_offset,
- HostMemoryVector* sink_size,
int depth) = 0;
+ HostMemoryVector* sink_size,
int depth,
+ bool layer_sliding_window =
false) = 0;
/*! \brief Copy the k position offset of applying RoPE for each sequence. */
virtual ffi::Tensor CopyKRoPEPosOffsetOnDepthAsync(HostMemoryVector* data,
int depth) = 0;
/*!
@@ -580,7 +583,8 @@ class PlainPagedKVCacheAuxDataManager : public
PagedKVCacheAuxDataManager {
explicit PlainPagedKVCacheAuxDataManager(int64_t reserved_num_seqs, int64_t
num_total_pages,
int64_t prefill_chunk_size,
DLDataType dtype_aux,
Device device, Device
preferred_host_device,
- TVMStreamHandle copy_stream)
+ TVMStreamHandle copy_stream,
+ bool layer_sliding_window = false)
: PagedKVCacheAuxDataManager(dtype_aux, device, preferred_host_device,
copy_stream),
cur_append_length_indptr_device_(
tvm::runtime::Empty({reserved_num_seqs + 1}, dtype_aux_, device)),
@@ -611,6 +615,14 @@ class PlainPagedKVCacheAuxDataManager : public
PagedKVCacheAuxDataManager {
tvm::runtime::Empty({num_total_pages}, dtype_aux_, device));
length_info_on_depths_device_.push_back(
tvm::runtime::Empty({3, reserved_num_seqs}, dtype_aux_, device));
+ if (layer_sliding_window) {
+ page_indptr_sliding_on_depths_device_.push_back(
+ tvm::runtime::Empty({reserved_num_seqs + 1}, dtype_aux_, device));
+ page_indices_sliding_on_depths_device_.push_back(
+ tvm::runtime::Empty({num_total_pages}, dtype_aux_, device));
+ length_info_sliding_on_depths_device_.push_back(
+ tvm::runtime::Empty({3, reserved_num_seqs}, dtype_aux_, device));
+ }
k_rope_pos_offset_on_depths_device_.push_back(
tvm::runtime::Empty({reserved_num_seqs}, dtype_aux_, device));
tree_attn_mask_device_.push_back(tvm::runtime::Empty(
@@ -628,15 +640,21 @@ class PlainPagedKVCacheAuxDataManager : public
PagedKVCacheAuxDataManager {
CopyVecDataToArray(view, data->data());
return view;
}
- ffi::Tensor CopyPageIndptrOnDepthAsync(HostMemoryVector* data, int depth)
final {
- ffi::Tensor view =
tvm::runtime::CreateView(page_indptr_on_depths_device_[depth],
-
{static_cast<int64_t>(data->size())}, dtype_aux_);
+ ffi::Tensor CopyPageIndptrOnDepthAsync(HostMemoryVector* data, int depth,
+ bool layer_sliding_window = false)
final {
+ const auto& storage = layer_sliding_window ?
page_indptr_sliding_on_depths_device_
+ : page_indptr_on_depths_device_;
+ ffi::Tensor view =
+ tvm::runtime::CreateView(storage[depth],
{static_cast<int64_t>(data->size())}, dtype_aux_);
CopyVecDataToArray(view, data->data());
return view;
}
- ffi::Tensor CopyPageIndicesOnDepthAsync(HostMemoryVector* data, int depth)
final {
- ffi::Tensor view =
tvm::runtime::CreateView(page_indices_on_depths_device_[depth],
-
{static_cast<int64_t>(data->size())}, dtype_aux_);
+ ffi::Tensor CopyPageIndicesOnDepthAsync(HostMemoryVector* data, int depth,
+ bool layer_sliding_window = false)
final {
+ const auto& storage = layer_sliding_window ?
page_indices_sliding_on_depths_device_
+ :
page_indices_on_depths_device_;
+ ffi::Tensor view =
+ tvm::runtime::CreateView(storage[depth],
{static_cast<int64_t>(data->size())}, dtype_aux_);
CopyVecDataToArray(view, data->data());
return view;
}
@@ -722,11 +740,13 @@ class PlainPagedKVCacheAuxDataManager : public
PagedKVCacheAuxDataManager {
ffi::Tensor CopyLengthInfoOnDepthAsync(HostMemoryVector* last_page_len,
HostMemoryVector*
sliding_window_offset,
- HostMemoryVector* sink_size, int
depth) final {
+ HostMemoryVector* sink_size, int
depth,
+ bool layer_sliding_window = false)
final {
int n_elem = last_page_len->size();
TVM_FFI_ICHECK_GT(n_elem, 0);
- ffi::Tensor view =
- tvm::runtime::CreateView(length_info_on_depths_device_[depth], {3,
n_elem}, dtype_aux_);
+ const auto& storage = layer_sliding_window ?
length_info_sliding_on_depths_device_
+ : length_info_on_depths_device_;
+ ffi::Tensor view = tvm::runtime::CreateView(storage[depth], {3, n_elem},
dtype_aux_);
ffi::Shape copy_shape{n_elem};
CopyVecDataToArray(view, last_page_len->data(), copy_shape);
CopyVecDataToArray(view, sliding_window_offset->data(), copy_shape,
@@ -804,6 +824,10 @@ class PlainPagedKVCacheAuxDataManager : public
PagedKVCacheAuxDataManager {
tvm::runtime::CopyFromTo(©_src, ©_dst, copy_stream_);
}
+ // Both page tables stay live until all global and sliding layers finish.
+ std::vector<ffi::Tensor> page_indptr_sliding_on_depths_device_;
+ std::vector<ffi::Tensor> page_indices_sliding_on_depths_device_;
+ std::vector<ffi::Tensor> length_info_sliding_on_depths_device_;
std::vector<ffi::Tensor> qo_indptr_on_depths_device_;
std::vector<ffi::Tensor> page_indptr_on_depths_device_;
std::vector<ffi::Tensor> page_indices_on_depths_device_;
@@ -865,10 +889,12 @@ class CachedPagedKVCacheAuxDataManager : public
PagedKVCacheAuxDataManager {
ffi::Tensor CopyQOIndptrOnDepthAsync(HostMemoryVector* data, int depth)
final {
return CopyAttnAuxVecToCache(data);
}
- ffi::Tensor CopyPageIndptrOnDepthAsync(HostMemoryVector* data, int depth)
final {
+ ffi::Tensor CopyPageIndptrOnDepthAsync(HostMemoryVector* data, int depth,
+ bool layer_sliding_window = false)
final {
return CopyAttnAuxVecToCache(data);
}
- ffi::Tensor CopyPageIndicesOnDepthAsync(HostMemoryVector* data, int depth)
final {
+ ffi::Tensor CopyPageIndicesOnDepthAsync(HostMemoryVector* data, int depth,
+ bool layer_sliding_window = false)
final {
return CopyAttnAuxVecToCache(data);
}
ffi::Tensor CopyLastPageLenOnDepthAsync(HostMemoryVector* data, int depth)
final {
@@ -914,7 +940,8 @@ class CachedPagedKVCacheAuxDataManager : public
PagedKVCacheAuxDataManager {
}
ffi::Tensor CopyLengthInfoOnDepthAsync(HostMemoryVector* last_page_len,
HostMemoryVector*
sliding_window_offset,
- HostMemoryVector* sink_size, int
depth) final {
+ HostMemoryVector* sink_size, int
depth,
+ bool layer_sliding_window = false)
final {
int64_t n_elem = last_page_len->size();
std::memcpy(merged_attn_aux_data_host_.data() + attn_aux_data_copy_offset_,
last_page_len->data(), n_elem * elem_byte_size_);
diff --git a/src/runtime/vm/paged_kv_cache.cc b/src/runtime/vm/paged_kv_cache.cc
index 68b826e52f..72e2b2e1f6 100644
--- a/src/runtime/vm/paged_kv_cache.cc
+++ b/src/runtime/vm/paged_kv_cache.cc
@@ -600,7 +600,7 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj
{
} else {
aux_data_manager_ = std::make_unique<PlainPagedKVCacheAuxDataManager>(
reserved_num_seqs, num_total_pages, prefill_chunk_size, dtype_aux_,
device,
- preferred_host_device, copy_stream_);
+ preferred_host_device, copy_stream_, support_layer_sliding_window_);
}
// Right now only the "normal" RoPE mode supports the RoPE extention
factors.
@@ -3171,7 +3171,7 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
qo_indptr_on_depths_host_[d].size());
page_indptr_sliding_window_on_depths_view_[d] =
aux_data_manager_->CopyPageIndptrOnDepthAsync(
- &page_indptr_sliding_window_on_depths_host_[d], d);
+ &page_indptr_sliding_window_on_depths_host_[d], d,
/*layer_sliding_window=*/true);
}
// 6. page_indices_sliding_window_on_depths
for (int d = 0; d < num_depths_; ++d) {
@@ -3179,7 +3179,7 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
page_indptr_sliding_window_on_depths_host_[d].back());
page_indices_sliding_window_on_depths_view_[d] =
aux_data_manager_->CopyPageIndicesOnDepthAsync(
- &page_indices_sliding_window_on_depths_host_[d], d);
+ &page_indices_sliding_window_on_depths_host_[d], d,
/*layer_sliding_window=*/true);
}
}
// 7. length_info_on_depths
@@ -3206,7 +3206,8 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
layer_sliding_window_length_info_on_depths_view_[d] =
aux_data_manager_->CopyLengthInfoOnDepthAsync(&last_page_len_on_depths_host_[d],
&sliding_window_offset_on_depths_host_[d],
-
&sink_size_on_depths_host_[d], d);
+
&sink_size_on_depths_host_[d], d,
+
/*layer_sliding_window=*/true);
}
}
// 6. k_rope_pos_offset_on_depths
diff --git a/tests/cpp/runtime/attn_aux_data_test.cc
b/tests/cpp/runtime/attn_aux_data_test.cc
new file mode 100644
index 0000000000..096daa517d
--- /dev/null
+++ b/tests/cpp/runtime/attn_aux_data_test.cc
@@ -0,0 +1,79 @@
+/*
+ * 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.
+ */
+
+#include <gtest/gtest.h>
+
+#include "../../../src/runtime/vm/attn_utils.h"
+
+namespace tvm {
+namespace runtime {
+namespace vm {
+namespace {
+
+HostMemoryVector IntVector(std::initializer_list<int32_t> values) {
+ HostMemoryVector result(values.size(), DLDataType{kDLInt, 32, 1},
Device{kDLCPU, 0});
+ for (int32_t value : values) result.push_back(value);
+ return result;
+}
+
+std::vector<int32_t> ReadInts(const ffi::Tensor& tensor) {
+ std::vector<int32_t> result(tensor.shape().Product());
+ tvm::runtime::CopyToBytes(tensor, result.data(), result.size() *
sizeof(int32_t));
+ return result;
+}
+
+TEST(PagedKVCacheAuxData, GlobalAndSlidingMetadataRemainIndependent) {
+ // Exercise the manager used by WebGPU and Metal on CPU. The normal CPU
+ // cache uses CachedPagedKVCacheAuxDataManager and would miss this aliasing.
+ Device cpu{kDLCPU, 0};
+ PlainPagedKVCacheAuxDataManager manager(2, 128, 512, DLDataType{kDLInt, 32,
1}, cpu, cpu, nullptr,
+ true);
+ auto full_indptr = IntVector({0, 4});
+ auto sliding_indptr = IntVector({0, 2});
+ auto full_indices = IntVector({10, 11, 12, 13});
+ auto sliding_indices = IntVector({12, 13});
+ auto last_page = IntVector({7});
+ auto full_offset = IntVector({0});
+ auto sliding_offset = IntVector({5});
+ auto sink = IntVector({0});
+ for (int iteration = 0; iteration < 2; ++iteration) {
+ manager.ResetAttnAuxDataCopy();
+ for (int depth = 0; depth < kPagedKVCacheMaxBlockDepth; ++depth) {
+ auto full_ip = manager.CopyPageIndptrOnDepthAsync(&full_indptr, depth);
+ auto full_ix = manager.CopyPageIndicesOnDepthAsync(&full_indices, depth);
+ auto full_len = manager.CopyLengthInfoOnDepthAsync(&last_page,
&full_offset, &sink, depth);
+ auto sliding_ip = manager.CopyPageIndptrOnDepthAsync(&sliding_indptr,
depth, true);
+ auto sliding_ix = manager.CopyPageIndicesOnDepthAsync(&sliding_indices,
depth, true);
+ auto sliding_len =
+ manager.CopyLengthInfoOnDepthAsync(&last_page, &sliding_offset,
&sink, depth, true);
+ manager.CommitAttnAuxDataCopy();
+ EXPECT_EQ(ReadInts(full_ip), (std::vector<int32_t>{0, 4}));
+ EXPECT_EQ(ReadInts(full_ix), (std::vector<int32_t>{10, 11, 12, 13}));
+ EXPECT_EQ(ReadInts(full_len), (std::vector<int32_t>{7, 0, 0}));
+ EXPECT_EQ(ReadInts(sliding_ip), (std::vector<int32_t>{0, 2}));
+ EXPECT_EQ(ReadInts(sliding_ix), (std::vector<int32_t>{12, 13}));
+ EXPECT_EQ(ReadInts(sliding_len), (std::vector<int32_t>{7, 5, 0}));
+ }
+ }
+}
+
+} // namespace
+} // namespace vm
+} // namespace runtime
+} // namespace tvm
