[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.

Author: akaashrp

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(&copy_src, &copy_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_);

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,15 +3171,15 @@ 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) {
         TVM_FFI_ICHECK_EQ(page_indices_sliding_window_on_depths_host_[d].size(),
                           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

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