diff --git a/modules/kv-state-cache/ds/kv_state_cache.cc b/modules/kv-state-cache/ds/kv_state_cache.cc
index 3437b223c..25ebbb193 100644
--- a/modules/kv-state-cache/ds/kv_state_cache.cc
+++ b/modules/kv-state-cache/ds/kv_state_cache.cc
@@ -83,7 +83,7 @@ KVStateCacheBuilder::KVStateCacheBuilder(Client& client, int dimension,
   std::shared_ptr<NodeData> rootTreeHeader = this->rootTree->GetRootNode();
   rootTreeHeader->treeData->data = treeData;
   rootTreeHeader->treeData->dataLength = sizeof(TreeData);
-  this->rootTree->SetSubtreeData(treeData, sizeof(TreeData));
+  this->rootTree->SetSubtreeData(treeData);
   LOG(INFO) << "set builder:" << builder
             << " to tree:" << this->rootTree->GetRootTree()->head;
   LOG(INFO) << "data:" << treeData
@@ -102,7 +102,7 @@ KVStateCacheBuilder::KVStateCacheBuilder(Client& client,
   std::set<void*> subTreeData = cache->rootTree->GetSubTreeDataSet();
 
   for (auto iter = subTreeData.begin(); iter != subTreeData.end(); ++iter) {
-    TreeData* treeData = (TreeData*) ((DataWrapper*) *iter)->data;
+    TreeData* treeData = (TreeData*) (*iter);
     LOG(INFO) << "tree data:" << treeData;
     VINEYARD_ASSERT(treeData->isPtr == false);
     LOG(INFO) << "id:" << treeData->builderObjectID;
@@ -135,13 +135,12 @@ KVStateCacheBlockBuilder* KVStateCacheBuilder::Split(
         kvStateCacheBlockBuilder->GetKeyStateBuilder();
     const std::shared_ptr<TensorBuilder<double>> valueStateTensorBuilder =
         kvStateCacheBlockBuilder->GetValueStateBuilder();
-    OffsetData* new_offset_data = new OffsetData();
+    OffsetData new_offset_data;
     childKVStateCacheBlockBuilder->Update(
         keyStateTensorBuilder->data() + index * this->dimension,
         valueStateTensorBuilder->data() + index * this->dimension,
-        this->dimension, new_offset_data);
-    nodeDataList[i]->nodeData->data = new_offset_data;
-    nodeDataList[i]->nodeData->dataLength = sizeof(OffsetData);
+        this->dimension, &new_offset_data);
+    data->offset = new_offset_data.offset;
     // Clear the bitmap.
     kvStateCacheBlockBuilder->DeleteKVCache(index);
   }
@@ -177,6 +176,13 @@ void KVStateCacheBuilder::Update(Client& client,
     Delete(evictedNodeData);
   }
 
+  // if (evictedNodeData->treeData != nullptr && evictedNodeData->nodeData !=
+  // nullptr) {
+  //   if (evictedNodeData->nodeData->data != nullptr) {
+  //     delete (TreeData*) evictedNodeData->nodeData->data;
+  //   }
+  // }
+
   // TBD
   // Use lock to protect the kv_state_cache_builder
   LOG(INFO) << "data:" << nodeData->treeData->data
@@ -204,7 +210,7 @@ void KVStateCacheBuilder::Update(Client& client,
 
     subTreeHeader->treeData->data = newTreeData;
     subTreeHeader->treeData->dataLength = sizeof(TreeData);
-    rootTree->SetSubtreeData(newTreeData, sizeof(TreeData));
+    rootTree->SetSubtreeData(newTreeData);
     LOG(INFO) << "block split success";
 
     // kv_state_cache_builder->UnLock();
@@ -257,6 +263,15 @@ void KVStateCacheBuilder::Delete(std::shared_ptr<NodeData> evictedNodeData) {
   kvStateCacheBlockBuilder->DeleteKVCache(data->offset);
   LOG(INFO) << "stage4";
   delete data;
+  // TBD
+  // Refactor this code. The data should be deleted by the RadixTree
+  // delete (DataWrapper*) evictedNodeData->nodeData;
+  LOG(INFO) << "tree data:" << evictedNodeData->treeData->data;
+  if (evictedNodeData->cleanTreeData) {
+    LOG(INFO) << "erase";
+    this->rootTree->GetSubTreeDataSet().erase(evictedNodeData->treeData->data);
+  }
+  evictedNodeData->RecycleSource();
 }
 
 void KVStateCacheBuilder::Merge(Client& client,
@@ -320,11 +335,10 @@ std::shared_ptr<Object> KVStateCacheBuilder::_Seal(Client& client) {
   // change the tree data from pointer to object id
 
   int count = 0;
-  LOG(INFO) << "count:" << count;
   std::set<void*> subTreeDataSet = rootTree->GetSubTreeDataSet();
   for (auto iter = subTreeDataSet.begin(); iter != subTreeDataSet.end();
        ++iter) {
-    TreeData* treeData = (TreeData*) ((DataWrapper*) *iter)->data;
+    TreeData* treeData = (TreeData*) (*iter);
     VINEYARD_ASSERT(treeData != nullptr);
     VINEYARD_ASSERT(treeData->isPtr == true);
 
@@ -357,12 +371,25 @@ std::shared_ptr<Object> KVStateCacheBuilder::_Seal(Client& client) {
 }
 
 KVStateCacheBuilder::~KVStateCacheBuilder() {
-  // TBD
-  // std::vector<std::shared_ptr<NodeData>> nodeDataList =
-  //     RadixTree::TraverseTreeWithoutSubTree(this->rootTree);
-  // for (size_t i = 0; i < nodeDataList.size(); i++) {
-  //   delete (OffsetData*) nodeDataList[i]->get_node()->get_data();
-  // }
+  LOG(INFO) << "KVStateCacheBuilder::~KVStateCacheBuilder";
+  // get all subtree data and node data
+  std::set<void*> subTreeDataSet = rootTree->GetSubTreeDataSet();
+  std::set<void*> nodeDataSet = rootTree->GetAllNodeData();
+  // 2. delete all subtree data and node data
+  for (auto iter = subTreeDataSet.begin(); iter != subTreeDataSet.end();
+       ++iter) {
+    TreeData* treeData = (TreeData*) (*iter);
+    if (treeData->isPtr == true) {
+      delete (KVStateCacheBlockBuilder*) treeData->kvStateCacheBlockBuilder;
+      delete treeData;
+    }
+  }
+  for (auto iter = nodeDataSet.begin(); iter != nodeDataSet.end(); ++iter) {
+    OffsetData* data = (OffsetData*) (*iter);
+    if (data != nullptr) {
+      delete data;
+    }
+  }
 }
 
 }  // namespace vineyard
diff --git a/modules/kv-state-cache/radix-tree/radix-tree.cc b/modules/kv-state-cache/radix-tree/radix-tree.cc
index c90063c55..1c2dda098 100644
--- a/modules/kv-state-cache/radix-tree/radix-tree.cc
+++ b/modules/kv-state-cache/radix-tree/radix-tree.cc
@@ -40,22 +40,23 @@ RadixTree::RadixTree(int cacheCapacity) {
   data->data = nullptr;
   data->dataLength = 0;
   dataNode->custom_data = data;
+  LOG(INFO) << "root data wrapper:" << data;
   dataNode->issubtree = true;
   this->rootToken = rootToken;
 }
 
 RadixTree::~RadixTree() {
-  // TBD
-  // raxFreeWithCallback(this->tree, [](raxNode *n) {
-  //   if (n->iskey && !n->isnull) {
-  //     nodeData* nodedata = (nodeData*) raxGetData(n);
-  //     delete nodedata;
-  //   }
-  //   if (n->issubtree && n->iscustomallocated && !n->iscustomnull) {
-  //     customData* customdata = (customData*) raxGetCustomData(n);
-  //     delete customdata;
-  //   }
-  // });
+  LOG(INFO) << "~RadixTree";
+  raxShow(this->tree);
+
+  raxNode* dataNode = raxFindAndReturnDataNode(this->tree, rootToken.data(),
+                                               rootToken.size(), NULL, false);
+  if (dataNode != nullptr) {
+    delete (DataWrapper*) dataNode->custom_data;
+    delete (DataWrapper*) raxGetData(dataNode);
+  }
+
+  raxFree(this->tree);
 }
 
 std::shared_ptr<NodeData> RadixTree::Insert(
@@ -154,15 +155,22 @@ void RadixTree::DeleteInternal(std::vector<int> tokens,
   DataWrapper* oldData;
   raxNode* subTreeNode;
   std::vector<int> pre;
-  // raxFindAndReturnDataNode(this->tree, deleteTokensArray,
-  // deleteTokensArrayLen,
-  //                          &subTreeNode, false);
+  raxNode* dataNode = raxFindAndReturnDataNode(
+      this->tree, deleteTokensArray, deleteTokensArrayLen, &subTreeNode, false);
+  bool nodeIsSubTree = false;
+  if (dataNode != nullptr && dataNode->issubtree) {
+    nodeIsSubTree = true;
+  }
   int retval = raxRemove(this->tree, deleteTokensArray, deleteTokensArrayLen,
-                         (void**) &oldData, &subTreeNode);
+                         (void**) &oldData);
   if (retval == 1) {
     evictedNode = std::make_shared<NodeData>(
         oldData, (DataWrapper*) subTreeNode->custom_data);
     nodeCount--;
+    if (nodeIsSubTree) {
+      // subTreeDataSet.erase(subTreeNode->custom_data);
+      evictedNode->cleanTreeData = true;
+    }
   } else {
     LOG(INFO) << "remove failed";
   }
@@ -224,6 +232,15 @@ std::string RadixTree::Serialize() {
 
     serializedStr += timestampOSS.str() + "|";
 
+    raxNode* node =
+        raxFindAndReturnDataNode(this->tree, tokenList[index].data(),
+                                 tokenList[index].size(), NULL, false);
+    uint32_t numNodes = node->numnodes;
+    std::ostringstream subTreeSizeOSS;
+    subTreeSizeOSS << std::hex << numNodes;
+
+    serializedStr += subTreeSizeOSS.str() + "|";
+
     // convert data to hex string
     char* bytes = (char*) ((DataWrapper*) dataList[index])->data;
     std::ostringstream dataOSS;
@@ -251,7 +268,7 @@ std::string RadixTree::Serialize() {
     char* bytes = (char*) ((DataWrapper*) subTreeDataList[index])->data;
     std::ostringstream dataOSS;
 
-    LOG(INFO) << "data lengtπh:"
+    LOG(INFO) << "data length:"
               << ((DataWrapper*) subTreeDataList[index])->dataLength;
     for (int i = 0; i < ((DataWrapper*) subTreeDataList[index])->dataLength;
          ++i) {
@@ -267,17 +284,21 @@ std::string RadixTree::Serialize() {
   // use ZSTD to compress the serialized string
   size_t srcSize = serializedStr.size();
   std::string compressedStr(srcSize, '\0');
-  int compressedSize = ZSTD_compress((void *)(compressedStr.c_str()), compressedStr.length(), 
-    serializedStr.c_str(), srcSize, 3);
+  int compressedSize =
+      ZSTD_compress((void*) (compressedStr.c_str()), compressedStr.length(),
+                    serializedStr.c_str(), srcSize, 3);
   if (ZSTD_isError(compressedSize)) {
-    LOG(ERROR) << "ZSTD compression failed: " << ZSTD_getErrorName(compressedSize);
+    LOG(ERROR) << "ZSTD compression failed: "
+               << ZSTD_getErrorName(compressedSize);
   }
   int cacheCapacity = this->cacheCapacity - 1;
 
-  std::string result = std::string((char*) &srcSize, sizeof(int)) +
-                       std::string((char*) &cacheCapacity, sizeof(int)) +
-                       compressedStr;
-  
+  std::string result =
+      std::string((char*) &srcSize, sizeof(int)) +
+      std::string((char*) &cacheCapacity, sizeof(int)) +
+      std::string((char*) &(this->tree->head->numnodes), sizeof(uint32_t)) +
+      compressedStr;
+
   return result;
 }
 
@@ -288,11 +309,15 @@ std::shared_ptr<RadixTree> RadixTree::Deserialize(std::string data) {
   data.erase(0, sizeof(int));
   int cacheCapacity = *(int*) data.c_str();
   data.erase(0, sizeof(int));
+  int rootNumNodes = *(uint32_t*) data.c_str();
+  data.erase(0, sizeof(uint32_t));
   std::string decompressedStr(srcSize, '\0');
-  int decompressedSize = ZSTD_decompress((void *)(decompressedStr.c_str()), decompressedStr.size(),
-   data.c_str(), srcSize);
+  int decompressedSize =
+      ZSTD_decompress((void*) (decompressedStr.c_str()), decompressedStr.size(),
+                      data.c_str(), srcSize);
   if (ZSTD_isError(decompressedSize)) {
-    LOG(ERROR) << "ZSTD decompression failed: " << ZSTD_getErrorName(decompressedSize);
+    LOG(ERROR) << "ZSTD decompression failed: "
+               << ZSTD_getErrorName(decompressedSize);
   }
   data = decompressedStr.substr(0, decompressedSize);
 
@@ -303,6 +328,7 @@ std::shared_ptr<RadixTree> RadixTree::Deserialize(std::string data) {
   std::vector<std::vector<int>> subTreeTokenList;
   std::vector<void*> subTreeDataList;
   std::vector<int> subTreeDataSizeList;
+  std::vector<int> subTreeSizeList;
   std::istringstream iss(data);
   std::string line;
   bool isMainTree = true;
@@ -315,7 +341,7 @@ std::shared_ptr<RadixTree> RadixTree::Deserialize(std::string data) {
     }
     LOG(INFO) << "data line:" << line << std::endl;
     std::istringstream lineStream(line);
-    std::string tokenListPart, timestampPart, dataPart;
+    std::string tokenListPart, timestampPart, dataPart, subTreeSizePart;
 
     if (!std::getline(lineStream, tokenListPart, '|')) {
       throw std::runtime_error(
@@ -326,6 +352,10 @@ std::shared_ptr<RadixTree> RadixTree::Deserialize(std::string data) {
         throw std::runtime_error(
             "Invalid serialized string format in timestamp part.");
       }
+      if (!std::getline(lineStream, subTreeSizePart, '|')) {
+        throw std::runtime_error(
+            "Invalid serialized string format in sub tree size part.");
+      }
     }
     if (!std::getline(lineStream, dataPart)) {
       LOG(INFO) << "data length is 0";
@@ -345,6 +375,15 @@ std::shared_ptr<RadixTree> RadixTree::Deserialize(std::string data) {
         LOG(INFO) << "Invalid timestamp format.";
         throw std::runtime_error("Invalid timestamp format.");
       }
+
+      std::istringstream subTreeSizeStream(subTreeSizePart);
+      uint32_t subTreeSize;
+      if (!(subTreeSizeStream >> std::hex >> subTreeSize)) {
+        LOG(INFO) << "Invalid sub tree size format.";
+        throw std::runtime_error("Invalid sub tree size format.");
+      }
+      LOG(INFO) << "Deserialize sub tree size:" << subTreeSize;
+      subTreeSizeList.push_back(subTreeSize);
     }
 
     size_t dataSize = dataPart.length() /
@@ -425,6 +464,16 @@ std::shared_ptr<RadixTree> RadixTree::Deserialize(std::string data) {
     }
     dataNode->timestamp = timestampList[i];
   }
+
+  for (size_t i = 0; i < tokenList.size(); i++) {
+    raxNode* node = raxFindAndReturnDataNode(
+        radixTree->tree, tokenList[i].data(), tokenList[i].size(), NULL, false);
+    LOG(INFO) << "node:" << node << " sub tree node num:" << subTreeSizeList[i];
+    node->numnodes = subTreeSizeList[i];
+  }
+  radixTree->tree->head->numnodes = rootNumNodes;
+  raxShow(radixTree->tree);
+
   LOG(INFO) << "start to insert sub tree token list" << std::endl;
   for (size_t i = 0; i < subTreeTokenList.size(); i++) {
     for (size_t j = 0; j < subTreeTokenList[i].size(); j++) {
@@ -449,20 +498,18 @@ std::shared_ptr<RadixTree> RadixTree::Deserialize(std::string data) {
     node->issubtree = true;
     raxSetCustomData(node, data);
 
-    // TBD
-    // refactor this code.
-    radixTree->subTreeDataSet.insert(data);
+    radixTree->subTreeDataSet.insert(subTreeDataList[i]);
   }
   LOG(INFO) << "Deserialize success";
+  raxShow(radixTree->tree);
   return radixTree;
 }
 
 std::vector<std::shared_ptr<NodeData>> RadixTree::SplitInternal(
     std::vector<int> tokens, std::shared_ptr<NodeData>& header) {
   std::vector<int> rootToken;
-  DataWrapper* dummyData = new DataWrapper();
   raxNode* subTreeRootNode =
-      raxSplit(this->tree, tokens.data(), tokens.size(), dummyData, rootToken);
+      raxSplit(this->tree, tokens.data(), tokens.size(), rootToken);
 
   raxShow(this->tree);
   subTreeRootNode->issubtree = true;
@@ -496,12 +543,9 @@ std::vector<std::shared_ptr<NodeData>> RadixTree::TraverseTreeWithoutSubTree(
   return nodes;
 }
 
-void RadixTree::SetSubtreeData(void* data, int dataLength) {
-  LOG(INFO) << "set subtree data";
-  DataWrapper* dataWrapper = new DataWrapper();
-  dataWrapper->data = data;
-  dataWrapper->dataLength = dataLength;
-  subTreeDataSet.insert(dataWrapper);
+void RadixTree::SetSubtreeData(void* data) {
+  LOG(INFO) << "set subtree data:" << data;
+  subTreeDataSet.insert(data);
 }
 
 std::shared_ptr<NodeData> RadixTree::GetRootNode() {
@@ -531,4 +575,19 @@ void RadixTree::MergeTree(std::shared_ptr<RadixTree> tree_1,
     std::vector<int> tmp(vec.begin() + 1, vec.end());
     insert_tokens.insert(tmp);
   }
+}
+
+std::set<void*> RadixTree::GetAllNodeData() {
+  raxIterator iter;
+  raxStart(&iter, this->tree);
+  raxSeek(&iter, "^", NULL, 0);
+  std::set<void*> nodeDataSet;
+  while (raxNext(&iter)) {
+    raxNode* node = iter.node;
+    if (node->isnull) {
+      continue;
+    }
+    nodeDataSet.insert(((DataWrapper*) raxGetData(node))->data);
+  }
+  return nodeDataSet;
 }
\ No newline at end of file
diff --git a/modules/kv-state-cache/radix-tree/radix-tree.h b/modules/kv-state-cache/radix-tree/radix-tree.h
index ce2d441c0..139f64e3c 100644
--- a/modules/kv-state-cache/radix-tree/radix-tree.h
+++ b/modules/kv-state-cache/radix-tree/radix-tree.h
@@ -37,11 +37,21 @@ struct DataWrapper {
 struct NodeData {
   DataWrapper* nodeData;
   DataWrapper* treeData;
+  bool cleanTreeData = false;
 
   NodeData(DataWrapper* nodeData, DataWrapper* treeData) {
     this->nodeData = nodeData;
     this->treeData = treeData;
   }
+
+  void RecycleSource() {
+    if (this->nodeData != nullptr) {
+      delete this->nodeData;
+    }
+    if (cleanTreeData && this->treeData != nullptr) {
+      delete this->treeData;
+    }
+  }
 };
 
 class RadixTree : public std::enable_shared_from_this<RadixTree> {
@@ -87,7 +97,7 @@ class RadixTree : public std::enable_shared_from_this<RadixTree> {
   static std::vector<std::shared_ptr<NodeData>> TraverseTreeWithoutSubTree(
       raxNode* headNode);
 
-  void SetSubtreeData(void* data, int dataLength);
+  void SetSubtreeData(void* data);
 
   rax* GetRootTree() { return this->tree; }
 
@@ -101,6 +111,8 @@ class RadixTree : public std::enable_shared_from_this<RadixTree> {
                         std::shared_ptr<RadixTree> tree_2,
                         std::vector<std::vector<int>>& evicted_tokens,
                         std::set<std::vector<int>>& insert_tokens);
+
+  std::set<void*> GetAllNodeData();
 };
 
 #endif
diff --git a/modules/kv-state-cache/radix-tree/radix.cc b/modules/kv-state-cache/radix-tree/radix.cc
index 590d75130..1373fb43c 100644
--- a/modules/kv-state-cache/radix-tree/radix.cc
+++ b/modules/kv-state-cache/radix-tree/radix.cc
@@ -155,10 +155,14 @@ static inline void raxStackFree(raxStack *ts) {
 /* Add the number of nodes in the stack to each node. */
 void raxStackAddNumNodes(raxStack *stack, int num) {
     for (size_t i=0; i<stack->items; i++) {
-        raxNode *node = (raxNode *)stack->stack[i];
+        raxNode *node = (raxNode *)stack->stack[stack->items - i - 1];
         node->numnodes+=(num);
+        if (node->issubtree) {
+            break;
+        }
     }
 }
+
 /* ----------------------------------------------------------------------------
  * Radix tree implementation
  * --------------------------------------------------------------------------*/
@@ -1097,6 +1101,25 @@ raxNode *raxFindAndReturnDataNode(rax *rax, int *s, size_t len, raxNode** sub_tr
     return h;
 }
 
+// raxNode *raxSetSubTreeAndReturnDataNode(rax *rax, int *s, size_t len) {
+//     raxNode *h;
+
+//     raxStack ts;
+//     raxStackInit(&ts);
+//     //debugf("### Lookup: %.*s\n", (int)len, s);
+//     int splitpos = 0;
+//     size_t i = raxLowWalk(rax,s,len,&h,NULL,&splitpos,&ts,false);
+//     if (i != len || (h->iscompr && splitpos != 0) || !h->iskey)
+//         return NULL;
+
+//     if (h!= nullptr) {
+//         h->issubtree = true;
+//         raxStackAddNumNodes(&ts, -h->numnodes);
+//     }
+
+//     return h;
+// }
+
 int raxFindNode(rax *rax, int *s, size_t len, void **node) {
     raxNode *h;
     raxStack ts;
@@ -1227,7 +1250,7 @@ raxNode *raxRemoveChild(raxNode *parent, raxNode *child) {
 
 /* Remove the specified item. Returns 1 if the item was found and
  * deleted, 0 otherwise. */
-int raxRemove(rax *rax, int *s, size_t len, void **old, raxNode** sub_tree_node, bool set_timestamp) {
+int raxRemove(rax *rax, int *s, size_t len, void **old, bool set_timestamp) {
     raxNode *h;
     raxStack ts;
 
@@ -1240,14 +1263,6 @@ int raxRemove(rax *rax, int *s, size_t len, void **old, raxNode** sub_tree_node,
         raxStackFree(&ts);
         return 0;
     }
-    if (sub_tree_node != NULL) {
-        for (int i = ts.items - 1; i >= 0; i--) {
-            if (((raxNode *)ts.stack[i])->issubtree == true) {
-                *sub_tree_node = (raxNode *)ts.stack[i];
-                break;
-            }
-        }
-    }
 
 
     if (old) *old = raxGetData(h);
@@ -2137,6 +2152,11 @@ uint64_t raxSize(rax *rax) {
  *  [1,2] -> [1,2,3,4] -> []
  */
 
+struct datawrapper {
+    void *data;
+    int length;
+};
+
 /* The actual implementation of raxShow(). */
 void raxRecursiveShow(int level, int lpad, raxNode *n) {
     char s = n->iscompr ? '"' : '[';
@@ -2156,6 +2176,9 @@ void raxRecursiveShow(int level, int lpad, raxNode *n) {
         numchars += printf("=%p",raxGetData(n));
     }
     numchars += printf(" node:%p time:%ld, data:%p, is_sub_tree:%d", n, n->timestamp, n->custom_data, n->issubtree);
+    if (n->issubtree && n->custom_data != NULL) {
+        numchars += printf(" cus data:%p" , ((datawrapper *)(n->custom_data))->data);
+    }
 
     int numchildren = n->iscompr ? 1 : n->size;
     /* Note that 7 and 4 magic constants are the string length
@@ -2184,6 +2207,7 @@ void raxRecursiveShow(int level, int lpad, raxNode *n) {
 
 /* Show a tree, as outlined in the comment above. */
 void raxShow(rax *rax) {
+    printf("rax numnode:%lu\n", rax->numele);
     raxRecursiveShow(0,0,rax->head);
     putchar('\n');
 }
@@ -2295,7 +2319,7 @@ bool raxIsSubtree(raxNode *node) {
 * tree from the root node.
 * 
 */
-raxNode *raxSplit(rax *rax, int *s, size_t len, void *data, std::vector<int>& token) {
+raxNode *raxSplit(rax *rax, int *s, size_t len, std::vector<int>& token) {
     raxNode *childNode = NULL;
     raxNode *splitNode = NULL;
     raxStack stack = raxFindWithStack(rax, s, len);
@@ -2349,7 +2373,7 @@ raxNode *raxSplit(rax *rax, int *s, size_t len, void *data, std::vector<int>& to
 
     raxSetSubtree(splitNode);
 
-    raxStackAddNumNodes(&stack, -(int)(splitNode->numnodes)); 
+    raxStackAddNumNodes(&stack, -(int)(splitNode->numnodes));
     raxStackFree(&stack);
 
     return splitNode;
diff --git a/modules/kv-state-cache/radix-tree/radix.h b/modules/kv-state-cache/radix-tree/radix.h
index 25321aa72..57da727cc 100644
--- a/modules/kv-state-cache/radix-tree/radix.h
+++ b/modules/kv-state-cache/radix-tree/radix.h
@@ -111,6 +111,7 @@ typedef struct raxNode {
   uint32_t numnodes;      /* Number of the child nodes */
   uint32_t numele;        /* Number of elements inside this node. */
   uint64_t timestamp;    /* Timestamps of the node */
+  uint32_t sub_tree_size; /* Number of nodes in the sub tree */
   void *custom_data;
   /* Data layout is as follows:
    *
@@ -212,7 +213,7 @@ int raxInsert(rax *rax, int *s, size_t len, void *data, void **old, bool set_tim
 int raxTryInsert(rax* rax, int* s, size_t len, void* data, void** old);
 int raxInsertAndReturnDataNode(rax* rax, int* s, size_t len, void* data,
                                void** node, void** old);
-int raxRemove(rax* rax, int* s, size_t len, void** old, raxNode** sub_tree_node = NULL, bool set_timestamp = true);
+int raxRemove(rax* rax, int* s, size_t len, void** old, bool set_timestamp = true);
 void* raxFind(rax* rax, int* s, size_t len);
 raxNode* raxFindAndReturnDataNode(rax* rax, int* s, size_t len, raxNode** sub_tree_node = NULL, bool set_timestamp = true);
 void raxSetSubtree(raxNode *n);
@@ -238,7 +239,7 @@ void raxSetDebugMsg(int onoff);
 void raxTraverse(raxNode* rax,
                  std::vector<std::shared_ptr<raxNode>>& dataNodeList);
 void raxTraverseSubTree(raxNode* n, std::vector<raxNode*> &dataNodeList);
-raxNode *raxSplit(rax *rax, int *s, size_t len, void *data, std::vector<int>& key);
+raxNode *raxSplit(rax *rax, int *s, size_t len, std::vector<int>& key);
 void raxSerialize(rax* root, std::vector<std::vector<int>>& tokenList, std::vector<void*>& dataList, std::vector<uint64_t> &timestampsList,
                   std::vector<std::vector<int>> *subtreeList, std::vector<void*> *subtreeNodeList);
 
@@ -251,4 +252,5 @@ void raxFindLastRecentNode(raxNode *node, std::vector<int>& key);
 void mergeTree(rax* first_tree, rax* second_tree, std::vector<std::vector<int>>& evicted_tokens, std::set<std::vector<int>>& insert_tokens, int max_node);
 void testIteRax(rax *tree);
 raxNode* raxGetFirstChildPtr(raxNode* node);
+// raxNode *raxSetSubTreeAndReturnDataNode(rax *rax, int *s, size_t len);
 #endif
diff --git a/test/kv_state_cache_test.cc b/test/kv_state_cache_test.cc
index 5782589f7..a1f4ed212 100644
--- a/test/kv_state_cache_test.cc
+++ b/test/kv_state_cache_test.cc
@@ -102,30 +102,23 @@ int main() {
   std::vector<int> round_2_tokens = {1, 2,  3,  4,  5,  7, 8,
                                      9, 10, 11, 12, 13, 14};
   std::vector<int> round_3_tokens = {1, 2, 3, 9, 10, 11, 12, 13, 14};
-  // total 24 node
-  // tree 1 : 18 node
-  // tree 2 : 16 node
-  // std::vector<int> round_3_tokens = {1, 2, 3, 4, 5, 6, 7};
-  // std::vector<int> round_1_tokens = {1, 2};
-  // std::vector<int> round_2_tokens = {1, 3};
-  // std::vector<int> round_3_tokens = {1, 3, 4};
-  // std::vector<int> round_4_tokens = {1, 3, 5};
-  // std::vector<int> round_5_tokens = {1, 1};
+  std::vector<int> round_4_tokens = {1, 2, 3, 4, 5, 6};
+
   inference(round_1_tokens);
   inference(round_2_tokens);
   sleep(5);
 
   inference(round_1_tokens);
   inference(round_2_tokens);
-  // sleep(5);
   inference(round_3_tokens);
-  // inference(round_3_tokens);
-  // inference(round_3_tokens);
   // inference(round_4_tokens);
-  // inference(round_5_tokens);
   // sleep(5);
-  // inference(round_2_tokens);
-  // inference(round_1_tokens, true);
+  // Delete(std::vector<int>(round_4_tokens.begin(), round_4_tokens.begin() +
+  // 6)); Delete(std::vector<int>(round_4_tokens.begin(), round_4_tokens.begin()
+  // + 5)); Delete(std::vector<int>(round_4_tokens.begin(),
+  // round_4_tokens.begin() + 4));
+  // Delete(std::vector<int>(round_4_tokens.begin(), round_4_tokens.begin() +
+  // 3));
   while (1)
     ;
   return 0;
diff --git a/test/rax_diff_test.cc b/test/rax_diff_test.cc
index e86731d10..ce1dcf1bf 100644
--- a/test/rax_diff_test.cc
+++ b/test/rax_diff_test.cc
@@ -79,7 +79,7 @@ int main(int argc, char** argv) {
   for (size_t i = 0; i < evicted_tokens.size(); i++) {
     // void* tree_data;
     raxRemove(rt_1, evicted_tokens[i].data(), evicted_tokens[i].size(), NULL,
-              NULL, false);
+              false);
   }
 
   for (auto it = insert_tokens.begin(); it != insert_tokens.end(); it++) {