From 246c6e55e14cf36cda65aaf921535544f2e24b0a Mon Sep 17 00:00:00 2001
From: Jiacheng Huang <huangjiacheng0709@outlook.com>
Date: Thu, 6 Aug 2026 20:41:48 +0800
Subject: [PATCH] refactor: migrate paged caching to canonical InfiniOps API

---
 .../paged_caching/paged_caching_infiniops.cc  | 23 +++++++++++++++----
 submodules/InfiniOps                          |  2 +-
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc b/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc
index 664472904..0c5c6ece4 100644
--- a/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc
+++ b/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc
@@ -1,23 +1,33 @@
 #include "infinicore/ops/paged_caching.hpp"
 
+#include <string>
+
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/paged_caching_infinilm.h"
+#include "base/reshape_and_cache_flash.h"
 
 namespace infinicore::op::paged_caching_impl::infiniops {
 namespace {
 using TensorMeta = ::infinicore::op::infiniops::TensorMeta;
 struct PlannedMeta {
-    TensorMeta k_cache, v_cache, k, v, slot_mapping;
-    graph::GraphTensor k_cache_tensor, v_cache_tensor, k_tensor, v_tensor, slot_mapping_tensor;
+    TensorMeta k, v, slot_mapping, scale, k_cache, v_cache;
+    graph::GraphTensor k_tensor, v_tensor, slot_mapping_tensor, scale_tensor, k_cache_tensor, v_cache_tensor;
 };
 } // namespace
 
 void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) {
     INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(k_cache->device().getType()));
     INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, slot_mapping);
-    return new PlannedMeta{TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(k), TensorMeta(v), TensorMeta(slot_mapping), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(slot_mapping)};
+
+    // The "auto" cache path ignores scales, but the canonical API requires them.
+    auto scale = Tensor::empty({1}, DataType::F32, k_cache->device());
+    auto k_cache_view = k_cache->permute({0, 2, 1, 3});
+    auto v_cache_view = v_cache->permute({0, 2, 1, 3});
+
+    return new PlannedMeta{
+        TensorMeta(k), TensorMeta(v), TensorMeta(slot_mapping), TensorMeta(scale), TensorMeta(k_cache_view), TensorMeta(v_cache_view),
+        graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(slot_mapping), graph::GraphTensor(scale), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache)};
 }
 
 void run(void *planned_meta) {
@@ -25,12 +35,15 @@ void run(void *planned_meta) {
     infini::ops::Handle handle;
     handle.set_stream(context::getStream());
     infini::ops::Config config;
-    infini::ops::PagedCachingInfinilm::Call(
+    infini::ops::ReshapeAndCacheFlash::Call(
         handle,
         config,
         planned->k.tensor(planned->k_tensor),
         planned->v.tensor(planned->v_tensor),
         planned->slot_mapping.tensor(planned->slot_mapping_tensor),
+        planned->scale.tensor(planned->scale_tensor),
+        planned->scale.tensor(planned->scale_tensor),
+        std::string{"auto"},
         planned->k_cache.tensor(planned->k_cache_tensor),
         planned->v_cache.tensor(planned->v_cache_tensor));
 }
diff --git a/submodules/InfiniOps b/submodules/InfiniOps
index 1b9c37659..e733e325b 160000
--- a/submodules/InfiniOps
+++ b/submodules/InfiniOps
@@ -1 +1 @@
-Subproject commit 1b9c37659211eeaeedb48b9691f79e4c9e4412f9
+Subproject commit e733e325b4af776261419e4dc7e1f9386c05fa71