From 86f60c345be3efafbe9dbeb39008cdf6caac11e3 Mon Sep 17 00:00:00 2001
From: Li Zhang <lzhang329@gmail.com>
Date: Thu, 2 Nov 2023 07:44:18 +0000
Subject: [PATCH] fix block verification

---
 src/turbomind/models/llama/BlockManager.cc    |  2 +-
 src/turbomind/models/llama/SequenceManager.cc | 24 ++++++++++---------
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/turbomind/models/llama/BlockManager.cc b/src/turbomind/models/llama/BlockManager.cc
index d04fd604b0..6ce8497eae 100644
--- a/src/turbomind/models/llama/BlockManager.cc
+++ b/src/turbomind/models/llama/BlockManager.cc
@@ -216,7 +216,7 @@ int BlockManager::Lock(const std::vector<const Block*>& bs)
 
     Move(cached_ids_, idxs, active_ids_);
 
-    dbg(cached_ids_, active_ids_);
+    // dbg(cached_ids_, active_ids_);
 
     return idxs.size();
 }
diff --git a/src/turbomind/models/llama/SequenceManager.cc b/src/turbomind/models/llama/SequenceManager.cc
index 6c2778daa1..c4d97ff539 100644
--- a/src/turbomind/models/llama/SequenceManager.cc
+++ b/src/turbomind/models/llama/SequenceManager.cc
@@ -97,9 +97,6 @@ bool SequenceManager::Erase(uint64_t id)
 
 void SequenceManager::VerifyAndLockCached(const Sequences& sequences)
 {
-    if (!need_verify_) {
-        return;
-    }
     std::vector<const Block*> blocks;
     for (const auto& p : sequences) {
         auto& seq = const_cast<Sequence&>(*p);
@@ -107,11 +104,13 @@ void SequenceManager::VerifyAndLockCached(const Sequences& sequences)
             continue;
         }
         FT_CHECK(seq.blocks.size() == seq.block_unique_ids.size());
-        for (int i = 0; i < seq.blocks.size(); ++i) {
-            if (seq.blocks[i]->unique_id != seq.block_unique_ids[i]) {
-                seq.blocks.resize(i);
-                seq.block_unique_ids.resize(i);
-                break;
+        if (need_verify_) {
+            for (int i = 0; i < seq.blocks.size(); ++i) {
+                if (seq.blocks[i]->unique_id != seq.block_unique_ids[i]) {
+                    seq.blocks.resize(i);
+                    seq.block_unique_ids.resize(i);
+                    break;
+                }
             }
         }
         blocks.insert(blocks.end(), seq.blocks.begin(), seq.blocks.end());
@@ -334,7 +333,7 @@ void SequenceManager::AssignAndActivate(const Sequences&                 sequenc
     for (int i = 0; i < sequences.size(); ++i) {
         auto& s     = const_cast<Sequence&>(*sequences[i]);
         auto  count = counts[i];
-        dbg(count);
+        // dbg(count);
         auto last = first + count;
         std::for_each(first, last, [&](const Block* b) {
             s.blocks.push_back(b);
@@ -417,8 +416,11 @@ auto SequenceManager::Materialize(Sequences                    sequences,
     }
 
     // allocate & assign blocks
-    if (schedule.allocate) {
-        auto blocks = block_manager_->Allocate(schedule.allocate);
+    {
+        std::vector<const Block*> blocks;
+        if (schedule.allocate) {
+            blocks = block_manager_->Allocate(schedule.allocate);
+        }
         AssignAndActivate(schedule.active, schedule.block_counts, blocks);
     }