diff --git a/modules/basic/ds/dataframe.cc b/modules/basic/ds/dataframe.cc index ebfdbdbf3..c0abe6edd 100644 --- a/modules/basic/ds/dataframe.cc +++ b/modules/basic/ds/dataframe.cc @@ -85,6 +85,12 @@ const std::shared_ptr DataFrame::AsBatch(bool copy) const { } else if (auto tensor = std::dynamic_pointer_cast>(df_col)) { num_rows = tensor->shape()[0]; + } else if (auto tensor = + std::dynamic_pointer_cast>(df_col)) { + num_rows = tensor->shape()[0]; + } else if (auto tensor = + std::dynamic_pointer_cast>(df_col)) { + num_rows = tensor->shape()[0]; } std::vector> buffer{ diff --git a/modules/llm-cache/ds/kv_state_cache.cc b/modules/llm-cache/ds/kv_state_cache.cc index 4c1615f82..3840d2b81 100644 --- a/modules/llm-cache/ds/kv_state_cache.cc +++ b/modules/llm-cache/ds/kv_state_cache.cc @@ -45,7 +45,9 @@ void KVStateCache::Resolve() { // 1. construct the radix tree this->rootTree = RadixTree::Deserialize( base64_decode(this->meta_.GetKeyValue("radix_tree"))); - // raxShow(this->rootTree->GetRootTree()); + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(this->rootTree->GetRootTree()); + } // 2. construct the kvStateCacheBlockBuilder list size_t numBlocks = this->meta_.GetKeyValue("numBlocks"); @@ -57,24 +59,24 @@ void KVStateCache::Resolve() { } // 3. construct the member field - this->dimension = this->meta_.GetKeyValue("dimension"); + this->tensorBytes = this->meta_.GetKeyValue("tensorBytes"); this->version = this->meta_.GetKeyValue("version"); this->layer = this->meta_.GetKeyValue("layer"); - VLOG(100) << "construct the member field success, with dimension:" - << this->dimension << " version:" << this->version + VLOG(100) << "construct the member field success, with tensorBytes:" + << this->tensorBytes << " version:" << this->version << " layer:" << this->layer; } KVStateCache::~KVStateCache() {} -KVStateCacheBuilder::KVStateCacheBuilder(Client& client, int dimension, +KVStateCacheBuilder::KVStateCacheBuilder(Client& client, int tensorBytes, int cacheCapacity, int layer, int blockSize) { - this->dimension = dimension; + this->tensorBytes = tensorBytes; this->version = 0; this->layer = layer; KVStateCacheBlockBuilder* builder = - new KVStateCacheBlockBuilder(client, this->dimension, layer, blockSize); + new KVStateCacheBlockBuilder(client, this->tensorBytes, layer, blockSize); this->rootTree = std::make_shared(cacheCapacity); @@ -90,7 +92,7 @@ KVStateCacheBuilder::KVStateCacheBuilder(Client& client, int dimension, KVStateCacheBuilder::KVStateCacheBuilder(Client& client, std::shared_ptr cache) { - this->dimension = cache->GetDimension(); + this->tensorBytes = cache->GetTensorBytes(); this->version = cache->GetVersion(); this->layer = cache->GetLayer(); // 1. create block builder from block @@ -114,11 +116,11 @@ KVStateCacheBuilder::KVStateCacheBuilder(Client& client, KVStateCacheBlockBuilder* KVStateCacheBuilder::Split( Client& client, KVStateCacheBlockBuilder* kvStateCacheBlockBuilder, - std::vector> nodeDataList) { + std::vector>& nodeDataList) { // Split the tree if the list of kvState is full. VINEYARD_ASSERT(nodeDataList.size() > 0); KVStateCacheBlockBuilder* childKVStateCacheBlockBuilder = - new KVStateCacheBlockBuilder(client, this->dimension, this->layer, + new KVStateCacheBlockBuilder(client, this->tensorBytes, this->layer, kvStateCacheBlockBuilder->GetBlockSize()); for (size_t i = 0; i < nodeDataList.size(); i++) { OffsetData* data = @@ -138,10 +140,9 @@ KVStateCacheBlockBuilder* KVStateCacheBuilder::Split( return childKVStateCacheBlockBuilder; } -void KVStateCacheBuilder::Update(Client& client, - const std::vector& tokenList, - int nextToken, - const KV_STATE_WITH_LAYER& kvState) { +void KVStateCacheBuilder::Update( + Client& client, const std::vector& tokenList, int nextToken, + const std::map>& kvState) { std::vector tokenListCopy = tokenList; tokenListCopy.push_back(nextToken); @@ -199,9 +200,9 @@ void KVStateCacheBuilder::Update(Client& client, << " bitmap:" << kvStateCacheBlockBuilder->GetBitmapStr(); } -int KVStateCacheBuilder::Query(Client& client, - const std::vector& tokenList, int token, - KV_STATE_WITH_LAYER& kvState) { +int KVStateCacheBuilder::Query( + Client& client, const std::vector& tokenList, int token, + std::map>& kvState) { std::vector tokenListCopy = tokenList; tokenListCopy.push_back(token); @@ -275,14 +276,14 @@ void KVStateCacheBuilder::Merge(Client& client, for (auto it = insertTokenList.begin(); it != insertTokenList.end(); ++it) { std::vector tokenList = std::vector((*it).begin(), (*it).end() - 1); - KV_STATE_WITH_LAYER kvState; + std::map> kvState; for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) { - K_STATE key_state; - V_STATE value_state; - key_state.data = malloc(this->dimension * sizeof(double)); - key_state.length = this->dimension * sizeof(double); - value_state.data = malloc(this->dimension * sizeof(double)); - value_state.length = this->dimension * sizeof(double); + LLMKV key_state; + LLMKV value_state; + key_state.data = malloc(this->tensorBytes); + key_state.length = this->tensorBytes; + value_state.data = malloc(this->tensorBytes); + value_state.length = this->tensorBytes; kvState.insert( std::make_pair(currentLayer, std::make_pair(key_state, value_state))); @@ -290,8 +291,8 @@ void KVStateCacheBuilder::Merge(Client& client, globalCacheBuilder->Query(client, tokenList, (*it).back(), kvState); this->Update(client, tokenList, (*it).back(), kvState); for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) { - K_STATE key_state = kvState[currentLayer].first; - V_STATE value_state = kvState[currentLayer].second; + LLMKV key_state = kvState[currentLayer].first; + LLMKV value_state = kvState[currentLayer].second; free(key_state.data); free(value_state.data); } @@ -309,7 +310,7 @@ std::shared_ptr KVStateCacheBuilder::_Seal(Client& client) { std::shared_ptr kvStateCache = std::make_shared(); // 1. store the member variables to cache object meta - kvStateCache->meta_.AddKeyValue("dimension", this->dimension); + kvStateCache->meta_.AddKeyValue("tensorBytes", this->tensorBytes); kvStateCache->meta_.AddKeyValue("version", this->version); kvStateCache->meta_.AddKeyValue("layer", this->layer); diff --git a/modules/llm-cache/ds/kv_state_cache.h b/modules/llm-cache/ds/kv_state_cache.h index 82e6a76c2..95b6d79df 100644 --- a/modules/llm-cache/ds/kv_state_cache.h +++ b/modules/llm-cache/ds/kv_state_cache.h @@ -15,6 +15,7 @@ limitations under the License. #include #include +#include #include #include "client/client.h" @@ -40,7 +41,7 @@ class KVStateCache : public vineyard::Registered { private: std::vector> kvStateCacheBlockList; std::shared_ptr rootTree; - int dimension; + int tensorBytes; int cacheCapacity; int layer; uint64_t version; @@ -56,11 +57,11 @@ class KVStateCache : public vineyard::Registered { void Resolve(); // for test - std::vector> GetKVStateCacheBlockList() { + std::vector>& GetKVStateCacheBlockList() { return this->kvStateCacheBlockList; } - int GetDimension() { return this->dimension; } + int GetTensorBytes() { return this->tensorBytes; } int GetCacheCapacity() { return this->cacheCapacity; } @@ -77,25 +78,26 @@ class KVStateCache : public vineyard::Registered { class KVStateCacheBuilder : public vineyard::ObjectBuilder { std::shared_ptr rootTree; - int dimension; + int tensorBytes; int layer; uint64_t version; public: - KVStateCacheBuilder(Client& client, int dimension, int cacheCapacity, + KVStateCacheBuilder(Client& client, int tensorBytes, int cacheCapacity, int layer, int blockSize = DEFAULT_BLOCK_SIZE); KVStateCacheBuilder(Client& client, std::shared_ptr cache); KVStateCacheBlockBuilder* Split( Client& client, KVStateCacheBlockBuilder* kvStateCacheBlockBuilder, - std::vector> nodeDataList); + std::vector>& nodeDataList); void Update(Client& client, const std::vector& token_list, - int next_token, const KV_STATE_WITH_LAYER& kv_state); + int next_token, + const std::map>& kv_state); int Query(Client& client, const std::vector& token_list, int token, - KV_STATE_WITH_LAYER& kv_state); + std::map>& kv_state); void Delete(std::shared_ptr evicted_node); @@ -109,7 +111,7 @@ class KVStateCacheBuilder : public vineyard::ObjectBuilder { std::shared_ptr _Seal(Client& client) override; - uint64_t GetDimension() { return this->dimension; } + uint64_t GetTensorBytes() { return this->tensorBytes; } std::shared_ptr GetRootTree() { return this->rootTree; } diff --git a/modules/llm-cache/ds/kv_state_cache_block.cc b/modules/llm-cache/ds/kv_state_cache_block.cc index 174771431..8f53f9357 100644 --- a/modules/llm-cache/ds/kv_state_cache_block.cc +++ b/modules/llm-cache/ds/kv_state_cache_block.cc @@ -60,10 +60,10 @@ void KVStateCacheBlock::Construct(const ObjectMeta& meta) { this->layer = this->meta_.GetKeyValue("layer"); for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) { this->keyStateTensorList.push_back( - std::dynamic_pointer_cast>(this->meta_.GetMember( + std::dynamic_pointer_cast>(this->meta_.GetMember( "keyStateTensorBuilder_" + std::to_string(currentLayer)))); this->valueStateTensorList.push_back( - std::dynamic_pointer_cast>(this->meta_.GetMember( + std::dynamic_pointer_cast>(this->meta_.GetMember( "valueStateTensorBuilder_" + std::to_string(currentLayer)))); } // 2. construct the member field @@ -74,27 +74,27 @@ void KVStateCacheBlock::Construct(const ObjectMeta& meta) { this->bitmap[i] = this->meta_.GetKeyValue("bitmap_" + std::to_string(i)); } - this->dimension = this->meta_.GetKeyValue("dimension"); + this->tensorBytes = this->meta_.GetKeyValue("tensorBytes"); this->blockSize = this->meta_.GetKeyValue("block_size"); } KVStateCacheBlock::~KVStateCacheBlock() { delete this->bitmap; } KVStateCacheBlockBuilder::KVStateCacheBlockBuilder(Client& client, - int dimension, int layer, + int tensorBytes, int layer, int blockSize) { this->blockSize = blockSize; this->bitmapSize = (blockSize + 63) / 64; this->bitmap = new uint64_t[this->bitmapSize]; memset(this->bitmap, UINT8_MAX, this->bitmapSize * sizeof(uint64_t)); - std::vector shape = {(int64_t)(blockSize), dimension}; + std::vector shape = {(int64_t)(blockSize), tensorBytes}; for (int i = 0; i < layer; i++) { this->keyStateTensorBuilderList.push_back( - std::make_shared>(client, shape)); + std::make_shared>(client, shape)); this->valueStateTensorBuilderList.push_back( - std::make_shared>(client, shape)); + std::make_shared>(client, shape)); } - this->dimension = dimension; + this->tensorBytes = tensorBytes; this->layer = layer; } @@ -108,37 +108,40 @@ KVStateCacheBlockBuilder::KVStateCacheBlockBuilder( for (int i = 0; i < this->bitmapSize; i++) { this->bitmap[i] = kvStateCacheBlock->bitmap[i]; } - this->dimension = kvStateCacheBlock->dimension; + this->tensorBytes = kvStateCacheBlock->tensorBytes; this->layer = kvStateCacheBlock->layer; - std::vector shape = {(int64_t)(blockSize), dimension}; + std::vector shape = {(int64_t)(blockSize), this->tensorBytes}; for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) { this->keyStateTensorBuilderList.push_back( - std::make_shared>(client, shape)); + std::make_shared>(client, shape)); this->valueStateTensorBuilderList.push_back( - std::make_shared>(client, shape)); + std::make_shared>(client, shape)); } for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) { memcpy(this->keyStateTensorBuilderList[currentLayer]->data(), kvStateCacheBlock->keyStateTensorList[currentLayer]->data(), - (int64_t)(blockSize) * this->dimension * sizeof(double)); + (int64_t)(blockSize) * this->tensorBytes); memcpy(this->valueStateTensorBuilderList[currentLayer]->data(), kvStateCacheBlock->valueStateTensorList[currentLayer]->data(), - (int64_t)(blockSize) * this->dimension * sizeof(double)); + (int64_t)(blockSize) * this->tensorBytes); } } // current we do not consider the layer. -int KVStateCacheBlockBuilder::Query(Client& client, int index, - KV_STATE_WITH_LAYER& kvState) { +int KVStateCacheBlockBuilder::Query( + Client& client, int index, + std::map>& kvState) { for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) { - memcpy((kvState.find(currentLayer)->second).first.data, - keyStateTensorBuilderList[currentLayer]->data() + index * dimension, - dimension * sizeof(double)); - memcpy( - (kvState.find(currentLayer)->second).second.data, - valueStateTensorBuilderList[currentLayer]->data() + index * dimension, - dimension * sizeof(double)); + LLMKV keyState = (kvState.find(currentLayer)->second).first; + LLMKV valueState = (kvState.find(currentLayer)->second).second; + keyState.data = + keyStateTensorBuilderList[currentLayer]->data() + index * tensorBytes; + keyState.length = tensorBytes; + valueState.data = + valueStateTensorBuilderList[currentLayer]->data() + index * tensorBytes; + valueState.length = tensorBytes; + kvState.emplace(currentLayer, std::make_pair(keyState, valueState)); } return 0; } @@ -164,23 +167,21 @@ bool KVStateCacheBlockBuilder::IsFull() { return true; } -void KVStateCacheBlockBuilder::Update(const KV_STATE_WITH_LAYER& kvState, - OffsetData* data) { +void KVStateCacheBlockBuilder::Update( + const std::map>& kvState, OffsetData* data) { int index = this->FindEmptySlot(); for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) { - K_STATE keyState = (kvState.find(currentLayer)->second).first; - V_STATE valueState = (kvState.find(currentLayer)->second).second; - VINEYARD_ASSERT(keyState.length == - (size_t) this->dimension * sizeof(double)); - VINEYARD_ASSERT(valueState.length == - (size_t) this->dimension * sizeof(double)); - - double* keyData = keyStateTensorBuilderList[currentLayer]->data(); - double* valueData = valueStateTensorBuilderList[currentLayer]->data(); - memcpy(keyData + index * this->dimension, keyState.data, - this->dimension * sizeof(double)); - memcpy(valueData + index * this->dimension, valueState.data, - this->dimension * sizeof(double)); + LLMKV keyState = (kvState.find(currentLayer)->second).first; + LLMKV valueState = (kvState.find(currentLayer)->second).second; + VINEYARD_ASSERT(keyState.length == (size_t) this->tensorBytes); + VINEYARD_ASSERT(valueState.length == (size_t) this->tensorBytes); + + uint8_t* keyData = keyStateTensorBuilderList[currentLayer]->data(); + uint8_t* valueData = valueStateTensorBuilderList[currentLayer]->data(); + memcpy(keyData + index * this->tensorBytes, keyState.data, + this->tensorBytes); + memcpy(valueData + index * this->tensorBytes, valueState.data, + this->tensorBytes); } data->offset = index; @@ -193,25 +194,26 @@ int16_t KVStateCacheBlockBuilder::Split(KVStateCacheBlockBuilder* child, VINEYARD_ASSERT(this->layer == child->layer); int childIndex = child->FindEmptySlot(); for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) { - std::shared_ptr> keyStateTensorBuilder = + std::shared_ptr> keyStateTensorBuilder = keyStateTensorBuilderList[currentLayer]; - std::shared_ptr> valueStateTensorBuilder = + std::shared_ptr> valueStateTensorBuilder = valueStateTensorBuilderList[currentLayer]; - std::shared_ptr> childKeyStateTensorBuilder = + std::shared_ptr> childKeyStateTensorBuilder = child->keyStateTensorBuilderList[currentLayer]; - std::shared_ptr> childValueStateTensorBuilder = + std::shared_ptr> childValueStateTensorBuilder = child->valueStateTensorBuilderList[currentLayer]; - double* keyState = keyStateTensorBuilder->data() + index * this->dimension; - double* valueState = - valueStateTensorBuilder->data() + index * this->dimension; - double* childKeyState = - childKeyStateTensorBuilder->data() + childIndex * this->dimension; - double* childValueState = - childValueStateTensorBuilder->data() + childIndex * this->dimension; - - memcpy(childKeyState, keyState, this->dimension * sizeof(double)); - memcpy(childValueState, valueState, this->dimension * sizeof(double)); + uint8_t* keyState = + keyStateTensorBuilder->data() + index * this->tensorBytes; + uint8_t* valueState = + valueStateTensorBuilder->data() + index * this->tensorBytes; + uint8_t* childKeyState = + childKeyStateTensorBuilder->data() + childIndex * this->tensorBytes; + uint8_t* childValueState = + childValueStateTensorBuilder->data() + childIndex * this->tensorBytes; + + memcpy(childKeyState, keyState, this->tensorBytes); + memcpy(childValueState, valueState, this->tensorBytes); } ACQUIRE_BIT_RESOURCE(child->bitmap[childIndex / 64], childIndex % 64); FREE_BIT_RESOURCE(this->bitmap[index / 64], index % 64); @@ -244,7 +246,7 @@ std::shared_ptr KVStateCacheBlockBuilder::_Seal(Client& client) { } kvStateCacheBlock->meta_.AddKeyValue("block_size", this->blockSize); - kvStateCacheBlock->meta_.AddKeyValue("dimension", this->dimension); + kvStateCacheBlock->meta_.AddKeyValue("tensorBytes", this->tensorBytes); kvStateCacheBlock->meta_.AddKeyValue("layer", this->layer); // 3. set the object type to meta kvStateCacheBlock->meta_.SetTypeName(type_name()); @@ -264,15 +266,16 @@ void KVStateCacheBlockBuilder::PrintKVStateCacheBlock() { LOG(INFO) << "layer:" << currentLayer; for (int i = 0; i < this->blockSize; i++) { LOG(INFO) << "index:" << i; + uint8_t* key_state_data = keyStateTensorBuilderList[currentLayer]->data(); + uint8_t* value_state_data = + valueStateTensorBuilderList[currentLayer]->data(); + // print the first tensorBytes bytes std::string keyState = ""; std::string valueState = ""; - for (int j = 0; j < this->dimension; j++) { - keyState += std::to_string((keyStateTensorBuilderList[currentLayer] - ->data())[i * dimension + j]) + - " "; - valueState += std::to_string((valueStateTensorBuilderList[currentLayer] - ->data())[i * dimension + j]) + - " "; + for (int j = 0; j < this->tensorBytes; j++) { + keyState += std::to_string(key_state_data[i * tensorBytes + j]) + " "; + valueState += + std::to_string(value_state_data[i * tensorBytes + j]) + " "; } LOG(INFO) << "keyState:" << keyState; LOG(INFO) << "valueState:" << valueState; diff --git a/modules/llm-cache/ds/kv_state_cache_block.h b/modules/llm-cache/ds/kv_state_cache_block.h index 5e0a7262b..f92f1a0f1 100644 --- a/modules/llm-cache/ds/kv_state_cache_block.h +++ b/modules/llm-cache/ds/kv_state_cache_block.h @@ -29,20 +29,11 @@ limitations under the License. #include "client/ds/i_object.h" #include "llm-cache/radix-tree/radix-tree.h" -struct State { +struct LLMKV { void* data; size_t length; }; -using K_STATE = State; -using V_STATE = State; - -using KV_STATE_WITH_LAYER = std::map>; -using LIST_KV_STATE_WITH_LAYER = - std::vector>>; -using KV_STATE = std::vector>; -using LIST_KV_STATE = std::vector>; - // Set the bit to 1, which means the resource is not being used #define FREE_BIT_RESOURCE(value, bit) ((value) |= (((uint64_t) 1) << (bit))) @@ -72,14 +63,14 @@ namespace vineyard { class KVStateCacheBlock : public vineyard::Registered { private: - std::vector>> keyStateTensorList; - std::vector>> valueStateTensorList; + std::vector>> keyStateTensorList; + std::vector>> valueStateTensorList; uint64_t* bitmap; int blockSize; int bitmapSize; ObjectID id; int layer; - int dimension; + int tensorBytes; public: static std::unique_ptr Create() __attribute__((used)) { @@ -91,25 +82,25 @@ class KVStateCacheBlock : public vineyard::Registered { std::string GetBitmapStr(); - uint64_t GetDimension() { return this->dimension; } + uint64_t GetTensorBytes() { return this->tensorBytes; } uint64_t* GetBitmap() { return this->bitmap; } int GetBlockSize() { return this->blockSize; } - std::shared_ptr> GetKeyTensor(int layer) { + std::shared_ptr> GetKeyTensor(int layer) { return this->keyStateTensorList[layer]; } - std::shared_ptr> GetValueTensor(int layer) { + std::shared_ptr> GetValueTensor(int layer) { return this->valueStateTensorList[layer]; } - std::vector>> GetKeyTensorList() { + std::vector>>& GetKeyTensorList() { return this->keyStateTensorList; } - std::vector>> GetValueTensorList() { + std::vector>>& GetValueTensorList() { return this->valueStateTensorList; } @@ -120,21 +111,22 @@ class KVStateCacheBlock : public vineyard::Registered { class KVStateCacheBlockBuilder : public ObjectBuilder { private: - std::vector>> keyStateTensorBuilderList; - std::vector>> + std::vector>> + keyStateTensorBuilderList; + std::vector>> valueStateTensorBuilderList; // TBD // support more than 64 kv-state cache slots uint64_t* bitmap; int blockSize; int bitmapSize; - int dimension; + int tensorBytes; int layer; int FindEmptySlot(); public: - KVStateCacheBlockBuilder(Client& client, int dimension, int layer, + KVStateCacheBlockBuilder(Client& client, int tensorBytes, int layer, int blockSize); KVStateCacheBlockBuilder( @@ -147,9 +139,10 @@ class KVStateCacheBlockBuilder : public ObjectBuilder { * @param kv_state The kv-state of the prompt. A LLM inference can contain * multiple kv-states for each layer. */ - void Update(const KV_STATE_WITH_LAYER& kv_state, OffsetData* data); + void Update(const std::map>& kv_state, + OffsetData* data); - void Update(double* keyState, double* valueState, uint64_t dataLength, + void Update(char* keyState, char* valueState, uint64_t dataLength, OffsetData* data); /** @@ -160,7 +153,8 @@ class KVStateCacheBlockBuilder : public ObjectBuilder { * @param kv_state The kv-state of the prompt returned by radix-tree. If the * kv-state is not found, the data of kv-state is invalid. */ - int Query(Client& client, int index, KV_STATE_WITH_LAYER& kv_state); + int Query(Client& client, int index, + std::map>& kv_state); bool IsFull(); @@ -170,20 +164,21 @@ class KVStateCacheBlockBuilder : public ObjectBuilder { int16_t Split(KVStateCacheBlockBuilder* child, int index); - const std::shared_ptr> GetKeyStateBuilder(int layer) { + const std::shared_ptr>& GetKeyStateBuilder(int layer) { return keyStateTensorBuilderList[layer]; } - const std::shared_ptr> GetValueStateBuilder(int layer) { + const std::shared_ptr>& GetValueStateBuilder( + int layer) { return valueStateTensorBuilderList[layer]; } - const std::vector>> + const std::vector>>& GetKeyStateBuilderList() { return keyStateTensorBuilderList; } - const std::vector>> + const std::vector>>& GetValueStateBuilderList() { return valueStateTensorBuilderList; } @@ -196,7 +191,7 @@ class KVStateCacheBlockBuilder : public ObjectBuilder { uint64_t* GetBitmap() { return this->bitmap; } - uint64_t GetDimension() { return this->dimension; } + uint64_t GetTensorBytes() { return this->tensorBytes; } int GetBlockSize() { return this->blockSize; } diff --git a/modules/llm-cache/ds/kv_state_cache_manager.cc b/modules/llm-cache/ds/kv_state_cache_manager.cc index b13770df4..c1cd75af0 100644 --- a/modules/llm-cache/ds/kv_state_cache_manager.cc +++ b/modules/llm-cache/ds/kv_state_cache_manager.cc @@ -25,7 +25,7 @@ limitations under the License. namespace vineyard { -KVStateCacheManager::KVStateCacheManager(int dimension, int cacheCapacity, +KVStateCacheManager::KVStateCacheManager(int tensorBytes, int cacheCapacity, int layer, int blockSize, int syncInterval, std::string socket) { this->syncInterval = syncInterval; @@ -61,7 +61,7 @@ KVStateCacheManager::KVStateCacheManager(int dimension, int cacheCapacity, // if failed, create a new cache object VLOG(100) << "failed to get the cache object, create a new one."; kvStateCacheBuilder = std::make_shared( - client, dimension, cacheCapacity, layer, blockSize); + client, tensorBytes, cacheCapacity, layer, blockSize); } // release the lock @@ -75,21 +75,21 @@ KVStateCacheManager::KVStateCacheManager(int dimension, int cacheCapacity, // use lease to prevent the deadlock if the client is down } -void KVStateCacheManager::UpdateInternal(const std::vector& tokenList, - int nextToken, - const KV_STATE_WITH_LAYER& kvState) { +void KVStateCacheManager::UpdateInternal( + const std::vector& tokenList, int nextToken, + const std::map>& kvState) { kvStateCacheBuilder->Update(client, tokenList, nextToken, kvState); } -int KVStateCacheManager::QueryInternal(const std::vector& tokenList, - int token, - KV_STATE_WITH_LAYER& kvState) { +int KVStateCacheManager::QueryInternal( + const std::vector& tokenList, int token, + std::map>& kvState) { return kvStateCacheBuilder->Query(client, tokenList, token, kvState); } -void KVStateCacheManager::Update(const std::vector& tokenList, - int nextToken, - const KV_STATE_WITH_LAYER& kvState) { +void KVStateCacheManager::Update( + const std::vector& tokenList, int nextToken, + const std::map>& kvState) { if (!syncMutex.try_lock()) { return; } @@ -99,8 +99,9 @@ void KVStateCacheManager::Update(const std::vector& tokenList, syncMutex.unlock(); } -void KVStateCacheManager::Update(const std::vector& tokenList, - const LIST_KV_STATE_WITH_LAYER& kvState) { +void KVStateCacheManager::Update( + const std::vector& tokenList, + const std::vector>>& kvState) { if (!syncMutex.try_lock()) { return; } @@ -114,8 +115,9 @@ void KVStateCacheManager::Update(const std::vector& tokenList, syncMutex.unlock(); } -int KVStateCacheManager::Query(const std::vector& tokenList, int token, - KV_STATE_WITH_LAYER& kvState) { +int KVStateCacheManager::Query( + const std::vector& tokenList, int token, + std::map>& kvState) { int result = -1; if (!syncMutex.try_lock()) { @@ -128,17 +130,25 @@ int KVStateCacheManager::Query(const std::vector& tokenList, int token, return result; } -int KVStateCacheManager::Query(const std::vector& tokenList, - LIST_KV_STATE_WITH_LAYER& listKVState) { +int KVStateCacheManager::Query( + const std::vector& tokenList, + std::vector>>& listKVState) { int result = -1; if (!syncMutex.try_lock()) { return result; } std::vector tokenListCopy; + std::map> kvState; for (size_t i = 0; i < tokenList.size(); i++) { - result = QueryInternal(tokenListCopy, tokenList[i], listKVState[i]); + result = QueryInternal(tokenListCopy, tokenList[i], kvState); + // if the result is -1, it means the token is not in the cache + if (result == -1) { + break; + } tokenListCopy.push_back(tokenList[i]); + listKVState.push_back(kvState); + kvState.clear(); } syncMutex.unlock(); @@ -158,11 +168,13 @@ KVStateCacheManager::~KVStateCacheManager() { } // This function is used for testing -void KVStateCacheManager::Delete(std::vector token) { +void KVStateCacheManager::Delete(std::vector& token) { std::shared_ptr evictedNode; kvStateCacheBuilder->GetRootTree()->Delete(token, evictedNode); kvStateCacheBuilder->Delete(evictedNode); - raxShow(kvStateCacheBuilder->GetRootTree()->tree); + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(kvStateCacheBuilder->GetRootTree()->tree); + } } void KVStateCacheManager::Sync() { diff --git a/modules/llm-cache/ds/kv_state_cache_manager.h b/modules/llm-cache/ds/kv_state_cache_manager.h index 408cac8a9..54da1e65c 100644 --- a/modules/llm-cache/ds/kv_state_cache_manager.h +++ b/modules/llm-cache/ds/kv_state_cache_manager.h @@ -14,9 +14,11 @@ limitations under the License. */ #include +#include #include #include #include +#include #include #include "llm-cache/ds/kv_state_cache.h" @@ -41,32 +43,33 @@ class KVStateCacheManager { public: KVStateCacheManager( - int dimension = 10, int cacheCapacity = 10, int layer = 1, + int tensorBytes = 80, int cacheCapacity = 10, int layer = 1, int blockSize = 5, int syncInterval = 3, std::string socket = std::string(getenv("VINEYARD_IPC_SOCKET"))); void Update(const std::vector& tokenList, int nextToken, - const KV_STATE_WITH_LAYER& kvState); + const std::map>& kvState); - void Update(const std::vector& tokenList, - const LIST_KV_STATE_WITH_LAYER& kvState); + void Update( + const std::vector& tokenList, + const std::vector>>& kvState); int Query(const std::vector& tokenList, int token, - KV_STATE_WITH_LAYER& kvState); + std::map>& kvState); int Query(const std::vector& tokenList, - LIST_KV_STATE_WITH_LAYER& listKVState); + std::vector>>& listKVState); ~KVStateCacheManager(); private: void UpdateInternal(const std::vector& tokenList, int nextToken, - const KV_STATE_WITH_LAYER& kvState); + const std::map>& kvState); int QueryInternal(const std::vector& tokenList, int token, - KV_STATE_WITH_LAYER& kvState); + std::map>& kvState); - void Delete(std::vector token); + void Delete(std::vector& token); static void SyncThreadFunc(KVStateCacheManager* manager); diff --git a/modules/llm-cache/radix-tree/radix-tree.cc b/modules/llm-cache/radix-tree/radix-tree.cc index 93e37a79e..2b1f1f286 100644 --- a/modules/llm-cache/radix-tree/radix-tree.cc +++ b/modules/llm-cache/radix-tree/radix-tree.cc @@ -33,7 +33,9 @@ RadixTree::RadixTree(int cacheCapacity) { std::vector rootToken = {INT32_MAX}; std::shared_ptr evictedNode; this->InsertInternal(rootToken, evictedNode); - // raxShow(this->tree); + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(this->tree); + } raxNode* dataNode = raxFindAndReturnDataNode(this->tree, rootToken.data(), rootToken.size(), NULL, false); DataWrapper* data = new DataWrapper(); @@ -47,7 +49,9 @@ RadixTree::RadixTree(int cacheCapacity) { RadixTree::~RadixTree() { VLOG(100) << "~RadixTree"; - // raxShow(this->tree); + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(this->tree); + } raxNode* dataNode = raxFindAndReturnDataNode(this->tree, rootToken.data(), rootToken.size(), NULL, false); @@ -60,30 +64,30 @@ RadixTree::~RadixTree() { } std::shared_ptr RadixTree::Insert( - std::vector tokens, std::shared_ptr& evictedNode) { + std::vector& tokens, std::shared_ptr& evictedNode) { tokens.insert(tokens.begin(), INT32_MAX); return InsertInternal(tokens, evictedNode); } -void RadixTree::Delete(std::vector tokens, +void RadixTree::Delete(std::vector& tokens, std::shared_ptr& evictedNode) { tokens.insert(tokens.begin(), INT32_MAX); DeleteInternal(tokens, evictedNode); } -std::shared_ptr RadixTree::Query(std::vector key) { +std::shared_ptr RadixTree::Query(std::vector& key) { key.insert(key.begin(), INT32_MAX); return QueryInternal(key); } std::vector> RadixTree::Split( - std::vector tokens, std::shared_ptr& header) { + std::vector& tokens, std::shared_ptr& header) { tokens.insert(tokens.begin(), INT32_MAX); return SplitInternal(tokens, header); } std::shared_ptr RadixTree::InsertInternal( - std::vector tokens, std::shared_ptr& evictedNode) { + const std::vector& tokens, std::shared_ptr& evictedNode) { // get the sub vector of the tokens std::vector rootToken = std::vector(tokens.begin(), tokens.end() - 1); @@ -92,7 +96,7 @@ std::shared_ptr RadixTree::InsertInternal( } // insert the token vector to the radix tree - int* insertTokensArray = tokens.data(); + const int* insertTokensArray = tokens.data(); size_t insertTokensArrayLen = tokens.size(); DataWrapper* dummyData = new DataWrapper(); DataWrapper* oldData; @@ -101,15 +105,16 @@ std::shared_ptr RadixTree::InsertInternal( this->tree, insertTokensArray, insertTokensArrayLen, dummyData, reinterpret_cast(&dataNode), reinterpret_cast(&oldData)); if (dataNode == NULL) { - throw std::runtime_error("Insert token list failed"); + LOG(ERROR) << "Insert token list failed"; return NULL; } if (retval == 1) { VLOG(100) << "node count++:" << this->nodeCount; nodeCount++; } - - // raxShow(this->tree); + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(this->tree); + } if (this->nodeCount > this->cacheCapacity) { VLOG(100) << "cache capacity is full, evict the last recent node"; VLOG(100) << "cache capacity:" << this->cacheCapacity @@ -143,12 +148,12 @@ std::shared_ptr RadixTree::InsertInternal( dummyData, reinterpret_cast(subTreeNode->custom_data)); } -void RadixTree::DeleteInternal(std::vector tokens, +void RadixTree::DeleteInternal(const std::vector& tokens, std::shared_ptr& evictedNode) { // remove the token vector from the radix tree // TBD // If the evicted node is the root node of sub tree, recycle the tree. - int* deleteTokensArray = tokens.data(); + const int* deleteTokensArray = tokens.data(); size_t deleteTokensArrayLen = tokens.size(); DataWrapper* oldData; @@ -174,9 +179,10 @@ void RadixTree::DeleteInternal(std::vector tokens, } } -std::shared_ptr RadixTree::QueryInternal(std::vector key) { +std::shared_ptr RadixTree::QueryInternal( + const std::vector& key) { VLOG(100) << "Query"; - int* tokens = key.data(); + const int* tokens = key.data(); size_t tokensLen = key.size(); if (this->tree == nullptr) { @@ -198,7 +204,9 @@ std::shared_ptr RadixTree::QueryInternal(std::vector key) { std::string RadixTree::Serialize() { VLOG(100) << "Serialize......"; - // raxShow(this->tree); + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(this->tree); + } std::vector> tokenList; std::vector dataList; std::vector timestampList; @@ -207,7 +215,9 @@ std::string RadixTree::Serialize() { raxSerialize(this->tree, tokenList, dataList, timestampList, &subTreeTokenList, &subTreeDataList); - // raxShow(this->tree); + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(this->tree); + } std::string serializedStr; if (tokenList.size() != dataList.size()) { @@ -447,7 +457,9 @@ std::shared_ptr RadixTree::Deserialize(std::string data) { std::make_shared(cacheCapacity); radixTree->nodeCount = tokenList.size(); - // raxShow(radixTree->tree); + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(radixTree->tree); + } for (size_t i = 0; i < tokenList.size(); i++) { std::string token_str = ""; for (size_t j = 0; j < tokenList[i].size(); j++) { @@ -480,7 +492,9 @@ std::shared_ptr RadixTree::Deserialize(std::string data) { node->numnodes = subTreeSizeList[i]; } radixTree->tree->head->numnodes = rootNumNodes; - // raxShow(radixTree->tree); + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(radixTree->tree); + } VLOG(100) << "start to insert sub tree token list" << std::endl; for (size_t i = 0; i < subTreeTokenList.size(); i++) { @@ -509,17 +523,18 @@ std::shared_ptr RadixTree::Deserialize(std::string data) { radixTree->SetSubtreeData(subTreeDataList[i]); } VLOG(100) << "Deserialize success"; - // raxShow(radixTree->tree); + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(radixTree->tree); + } return radixTree; } std::vector> RadixTree::SplitInternal( - std::vector tokens, std::shared_ptr& header) { + const std::vector& tokens, std::shared_ptr& header) { std::vector rootToken; raxNode* subTreeRootNode = raxSplit(this->tree, tokens.data(), tokens.size(), rootToken); - // raxShow(this->tree); subTreeRootNode->issubtree = true; DataWrapper* treeData = new DataWrapper(); treeData->data = nullptr; diff --git a/modules/llm-cache/radix-tree/radix-tree.h b/modules/llm-cache/radix-tree/radix-tree.h index 211b48f9c..485a6ad69 100644 --- a/modules/llm-cache/radix-tree/radix-tree.h +++ b/modules/llm-cache/radix-tree/radix-tree.h @@ -65,30 +65,30 @@ class RadixTree : public std::enable_shared_from_this { private: std::shared_ptr InsertInternal( - std::vector tokens, std::shared_ptr& evictedNode); + const std::vector& tokens, std::shared_ptr& evictedNode); - void DeleteInternal(std::vector tokens, + void DeleteInternal(const std::vector& tokens, std::shared_ptr& evictedNode); - std::shared_ptr QueryInternal(std::vector key); + std::shared_ptr QueryInternal(const std::vector& key); std::vector> SplitInternal( - std::vector tokens, std::shared_ptr& header); + const std::vector& tokens, std::shared_ptr& header); public: RadixTree(int cacheCapacity); // NOLINT(runtime/explicit) ~RadixTree(); - std::shared_ptr Insert(std::vector tokens, + std::shared_ptr Insert(std::vector& tokens, std::shared_ptr& evictedNode); - void Delete(std::vector tokens, std::shared_ptr& evictedNode); + void Delete(std::vector& tokens, std::shared_ptr& evictedNode); - std::shared_ptr Query(std::vector key); + std::shared_ptr Query(std::vector& key); std::vector> Split( - std::vector tokens, std::shared_ptr& header); + std::vector& tokens, std::shared_ptr& header); std::string Serialize(); diff --git a/modules/llm-cache/tests/kv_state_cache_benchmark_test.cc b/modules/llm-cache/tests/kv_state_cache_benchmark_test.cc index feb1166b4..ab5a986cd 100644 --- a/modules/llm-cache/tests/kv_state_cache_benchmark_test.cc +++ b/modules/llm-cache/tests/kv_state_cache_benchmark_test.cc @@ -27,16 +27,15 @@ limitations under the License. using namespace vineyard; // NOLINT(build/namespaces) -#define DIMENSION 100 -#define CAPACITY 1000 -#define LAYER 64 -#define BLOCK_SIZE 100 +constexpr int TENSORBYTES = 800; +constexpr int CAPACITY = 1000; +constexpr int LAYER = 64; +constexpr int BLOCK_SIZE = 100; KVStateCacheManager* manager; void init() { - manager = - new KVStateCacheManager(DIMENSION, CAPACITY, LAYER, DEFAULT_BLOCK_SIZE); + manager = new KVStateCacheManager(TENSORBYTES, CAPACITY, LAYER, BLOCK_SIZE); } std::vector generate_random_tokens(size_t max_length) { @@ -52,15 +51,15 @@ std::vector generate_random_tokens(size_t max_length) { return tokens; } -std::map> generate_kv_state(int token) { - std::map> kv_state; +std::map> generate_kv_state(int token) { + std::map> kv_state; for (int currentLayer = 0; currentLayer < LAYER; currentLayer++) { - K_STATE key_state; - V_STATE value_state; - key_state.data = malloc(DIMENSION * sizeof(double)); - key_state.length = DIMENSION * sizeof(double); - value_state.data = malloc(DIMENSION * sizeof(double)); - value_state.length = DIMENSION * sizeof(double); + LLMKV key_state; + LLMKV value_state; + key_state.data = malloc(TENSORBYTES); + key_state.length = TENSORBYTES; + value_state.data = malloc(TENSORBYTES); + value_state.length = TENSORBYTES; kv_state.insert( std::make_pair(currentLayer, std::make_pair(key_state, value_state))); @@ -71,7 +70,7 @@ std::map> generate_kv_state(int token) { // test the performance of Query and Update function void benchmark_inference(std::vector>& tokens) { LOG(INFO) << "inference for benchmark"; - std::map> kv_state; + std::map> kv_state; std::chrono::steady_clock::time_point start, end; double token_list_size = 0; diff --git a/modules/llm-cache/tests/kv_state_cache_multi_test.cc b/modules/llm-cache/tests/kv_state_cache_multi_test.cc index 52d4ae426..1850bf193 100644 --- a/modules/llm-cache/tests/kv_state_cache_multi_test.cc +++ b/modules/llm-cache/tests/kv_state_cache_multi_test.cc @@ -22,14 +22,7 @@ limitations under the License. #include "common/util/logging.h" -char process_name[] = "kv_state_cache_test"; -char arg_0[] = "-s"; -char token_sequence_1[] = "1"; -char token_sequence_2[] = "2"; -char token_sequence_3[] = "3"; -char token_sequence_4[] = "4"; - -const char* program = "./build/bin/kv_state_cache_test"; +constexpr char* program = "./build/bin/kv_state_cache_test"; pid_t create_subprocess(char* argv[]) { pid_t pid = fork(); @@ -46,6 +39,13 @@ pid_t create_subprocess(char* argv[]) { } int main(int argc, char** argv) { + char process_name[] = "kv_state_cache_test"; + char arg_0[] = "-s"; + char token_sequence_1[] = "1"; + char token_sequence_2[] = "2"; + char token_sequence_3[] = "3"; + char token_sequence_4[] = "4"; + std::string sockets[2]; std::string rpc_endpoint; for (int i = 1; i < argc; i++) { diff --git a/modules/llm-cache/tests/kv_state_cache_radix_tree_test.cc b/modules/llm-cache/tests/kv_state_cache_radix_tree_test.cc index d3f4ae780..4fc3b3860 100644 --- a/modules/llm-cache/tests/kv_state_cache_radix_tree_test.cc +++ b/modules/llm-cache/tests/kv_state_cache_radix_tree_test.cc @@ -40,15 +40,19 @@ void radix_tree_insert_test() { std::shared_ptr node_data; for (int i = 0; i < 10; i++) { tokens.push_back(i); - VINEYARD_ASSERT(radix_tree->Insert(tokens, node_data) != NULL); + std::vector tokens_copy = tokens; + VINEYARD_ASSERT(radix_tree->Insert(tokens_copy, node_data) != NULL); VINEYARD_ASSERT(node_data == NULL); } - + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(radix_tree->tree); + } /* insert new token and check whether the old token is evicted */ tokens.clear(); for (int i = 1; i < 10; i++) { tokens.push_back(i); - VINEYARD_ASSERT(radix_tree->Insert(tokens, node_data) != NULL); + std::vector tokens_copy = tokens; + VINEYARD_ASSERT(radix_tree->Insert(tokens_copy, node_data) != NULL); VINEYARD_ASSERT(node_data != NULL); } @@ -68,7 +72,8 @@ void radix_tree_delete_test() { std::shared_ptr node_data; for (int i = 0; i < 10; i++) { tokens.push_back(i); - VINEYARD_ASSERT(radix_tree->Insert(tokens, node_data) != NULL); + std::vector tokens_copy = tokens; + VINEYARD_ASSERT(radix_tree->Insert(tokens_copy, node_data) != NULL); VINEYARD_ASSERT(node_data == NULL); } @@ -100,7 +105,8 @@ void radix_tree_query_test() { std::shared_ptr node_data; for (int i = 0; i < 10; i++) { tokens.push_back(i); - VINEYARD_ASSERT(radix_tree->Insert(tokens, node_data) != NULL); + std::vector tokens_copy = tokens; + VINEYARD_ASSERT(radix_tree->Insert(tokens_copy, node_data) != NULL); VINEYARD_ASSERT(node_data == NULL); } @@ -127,7 +133,8 @@ void radix_tree_serialize_and_deserialize() { std::shared_ptr node_data; for (int i = 0; i < 10; i++) { tokens.push_back(i); - VINEYARD_ASSERT(radix_tree->Insert(tokens, node_data) != NULL); + std::vector tokens_copy = tokens; + VINEYARD_ASSERT(radix_tree->Insert(tokens_copy, node_data) != NULL); VINEYARD_ASSERT(node_data == NULL); } @@ -142,21 +149,22 @@ void radix_tree_serialize_and_deserialize() { tokens.clear(); for (int i = 0; i < 10; i++) { tokens.push_back(i); + std::vector tokens_copy = tokens; print_tokens(tokens); - VINEYARD_ASSERT(deserialized_radix_tree->Query(tokens) != NULL); + VINEYARD_ASSERT(deserialized_radix_tree->Query(tokens_copy) != NULL); } } void radix_tree_split() { std::shared_ptr radix_tree = std::make_shared(20); - raxShow(radix_tree->tree); /* insert a token list*/ std::vector tokens; std::shared_ptr node_data; for (int i = 0; i < 10; i++) { tokens.push_back(i); - VINEYARD_ASSERT(radix_tree->Insert(tokens, node_data) != NULL); + std::vector tokens_copy = tokens; + VINEYARD_ASSERT(radix_tree->Insert(tokens_copy, node_data) != NULL); VINEYARD_ASSERT(node_data == NULL); } diff --git a/modules/llm-cache/tests/kv_state_cache_test.cc b/modules/llm-cache/tests/kv_state_cache_test.cc index e2d1e98ac..4776c264d 100644 --- a/modules/llm-cache/tests/kv_state_cache_test.cc +++ b/modules/llm-cache/tests/kv_state_cache_test.cc @@ -24,27 +24,18 @@ limitations under the License. using namespace vineyard; // NOLINT(build/namespaces) -int dimension = 10; +int tensorBytes = 80; int capacity = 20; int layer = 3; int block_size = 5; -std::vector round_1_tokens = { - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, - 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, - 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, - 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69}; -std::vector round_2_tokens = {1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14}; -std::vector round_3_tokens = {1, 2, 3, 9, 10, 11, 12, 13, 14}; -std::vector round_4_tokens = {1, 2, 3, 4, 5, 6}; - std::vector> tokens_list; KVStateCacheManager* kv_state_cache_manager; -void init(int dimension, int capacity, int layer, int block_size, +void init(int tensorBytes, int capacity, int layer, int block_size, std::string socket) { - kv_state_cache_manager = new KVStateCacheManager(dimension, capacity, layer, + kv_state_cache_manager = new KVStateCacheManager(tensorBytes, capacity, layer, block_size, 3, socket); } @@ -57,21 +48,19 @@ void print_current_tokens(const std::vector& prefix, int next_token) { LOG(INFO) << "Current tokens: " + tokens_str; } -void print_kv_state( - const std::map>& kv_state) { +void print_kv_state(const std::map>& kv_state) { LOG(INFO) << "kv_state: "; for (auto iter = kv_state.begin(); iter != kv_state.end(); ++iter) { + uint8_t* key_state_data = + reinterpret_cast(iter->second.first.data); + uint8_t* value_state_data = + reinterpret_cast(iter->second.second.data); + // print the first tensorBytes bytes std::string key_state_str = ""; std::string value_state_str = ""; - for (int i = 0; i < dimension; ++i) { - key_state_str += - std::to_string( - (reinterpret_cast(iter->second.first.data))[i]) + - " "; - value_state_str += - std::to_string( - (reinterpret_cast(iter->second.second.data))[i]) + - " "; + for (int j = 0; j < tensorBytes; j++) { + key_state_str += std::to_string(key_state_data[j]) + " "; + value_state_str += std::to_string(value_state_data[j]) + " "; } LOG(INFO) << "layer " << iter->first << ":"; LOG(INFO) << "key_state: " << key_state_str; @@ -81,69 +70,59 @@ void print_kv_state( } // we do not consider the layer. -std::map> generate_kv_state() { - std::map> kv_state; +std::map> generate_kv_state(int token) { + std::map> kv_state; for (int currentLayer = 0; currentLayer < layer; currentLayer++) { - K_STATE key_state; - V_STATE value_state; - key_state.data = malloc(dimension * sizeof(double)); - value_state.data = malloc(dimension * sizeof(double)); + LLMKV key_state; + LLMKV value_state; + key_state.data = malloc(tensorBytes); + value_state.data = malloc(tensorBytes); - key_state.length = dimension * sizeof(double); - value_state.length = dimension * sizeof(double); + key_state.length = tensorBytes; + value_state.length = tensorBytes; - kv_state.insert( - std::make_pair(currentLayer, std::make_pair(key_state, value_state))); - } - return kv_state; -} - -void update_kv_state(std::map>& kvState, - int token) { - for (int currentLayer = 0; currentLayer < layer; currentLayer++) { - K_STATE key_state = kvState[currentLayer].first; - V_STATE value_state = kvState[currentLayer].second; - for (int i = 0; i < dimension; ++i) { - (reinterpret_cast(key_state.data))[i] = - (static_cast(token)) / dimension * (i + 1) + + for (int i = 0; i < tensorBytes; ++i) { + (reinterpret_cast(key_state.data))[i] = + (static_cast(token)) / tensorBytes * (i + 1) + currentLayer * 10; - (reinterpret_cast(value_state.data))[i] = - (static_cast(token)) / dimension * (i + 1) * 2 + + (reinterpret_cast(value_state.data))[i] = + (static_cast(token)) / tensorBytes * (i + 1) * 2 + currentLayer * 10; } + kv_state.insert( + std::make_pair(currentLayer, std::make_pair(key_state, value_state))); } + return kv_state; } -void check_kv_state(const std::map>& kv_state, +void check_kv_state(const std::map>& kv_state, int& token) { VINEYARD_ASSERT(kv_state.size() == (size_t) layer); for (auto iter = kv_state.begin(); iter != kv_state.end(); ++iter) { - VINEYARD_ASSERT(iter->second.first.length == - (size_t) dimension * sizeof(double)); - VINEYARD_ASSERT(iter->second.second.length == - (size_t) dimension * sizeof(double)); - for (int i = 0; i < dimension; ++i) { - if ((reinterpret_cast(iter->second.first.data))[i] != - (static_cast(token)) / dimension * (i + 1) + + VINEYARD_ASSERT(iter->second.first.length == (size_t) tensorBytes); + VINEYARD_ASSERT(iter->second.second.length == (size_t) tensorBytes); + for (int i = 0; i < tensorBytes; ++i) { + if ((reinterpret_cast(iter->second.first.data))[i] != + (static_cast(token)) / tensorBytes * (i + 1) + iter->first * 10) { - LOG(INFO) << "token:" << token << " dimension" << dimension + LOG(INFO) << "token:" << token << " tensorBytes" << tensorBytes << " layer:" << iter->first; LOG(INFO) << "key_state[" << i << "]: " - << (reinterpret_cast(iter->second.first.data))[i] + << (reinterpret_cast(iter->second.first.data))[i] << ". But is should be " - << (static_cast(token)) / dimension * (i + 1) + + << (static_cast(token)) / tensorBytes * (i + 1) + iter->first * 10; throw std::runtime_error("key_state error!"); } - if ((reinterpret_cast(iter->second.second.data))[i] != - (static_cast(token)) / dimension * (i + 1) * 2 + + if (reinterpret_cast(iter->second.second.data)[i] != + (static_cast(token)) / tensorBytes * (i + 1) * 2 + iter->first * 10) { - LOG(INFO) << "token:" << token << " dimension" << dimension + LOG(INFO) << "token:" << token << " tensorBytes" << tensorBytes << " layer:" << iter->first; LOG(INFO) << "value_state[" << i << "]: " - << (reinterpret_cast(iter->second.second.data))[i] + << (reinterpret_cast(iter->second.second.data))[i] << ". But is should be " - << (static_cast(token)) / dimension * (i + 1) * 2 + + << (static_cast(token)) / tensorBytes * (i + 1) * 2 + iter->first * 10; throw std::runtime_error("value_state error!"); } @@ -153,16 +132,16 @@ void check_kv_state(const std::map>& kv_state, void inference(std::vector tokens, bool block = false) { std::vector inference_tokens; - std::map> kv_state; - kv_state = generate_kv_state(); + std::map> kv_state; for (size_t i = 0; i < tokens.size(); ++i) { + kv_state.clear(); int result = kv_state_cache_manager->Query(inference_tokens, tokens[i], kv_state); if (result != 0) { LOG(INFO) << "Can not find the kv_state from cache:"; print_current_tokens(inference_tokens, tokens[i]); LOG(INFO) << "Generate the kv_state and update the cache."; - update_kv_state(kv_state, tokens[i]); + kv_state = generate_kv_state(tokens[i]); print_kv_state(kv_state); kv_state_cache_manager->Update(inference_tokens, tokens[i], kv_state); } else { @@ -176,6 +155,15 @@ void inference(std::vector tokens, bool block = false) { } int main(int argc, char** argv) { + std::vector round_1_tokens = { + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, + 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, + 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, + 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69}; + std::vector round_2_tokens = {1, 2, 3, 4, 5, 7, 8, + 9, 10, 11, 12, 13, 14}; + std::vector round_3_tokens = {1, 2, 3, 9, 10, 11, 12, 13, 14}; + std::vector round_4_tokens = {1, 2, 3, 4, 5, 6}; if (argc < 2) { printf("usage ./kv_state_cache_test "); return 1; @@ -184,7 +172,7 @@ int main(int argc, char** argv) { for (int i = 2; i < argc; i++) { if (strcmp(argv[i], "-d") == 0) { - dimension = atoi(argv[i + 1]); + tensorBytes = atoi(argv[i + 1]); } else if (strcmp(argv[i], "-c") == 0) { capacity = atoi(argv[i + 1]); } else if (strcmp(argv[i], "-l") == 0) { @@ -209,11 +197,11 @@ int main(int argc, char** argv) { } } - LOG(INFO) << "Test KVStateCache with dimension: " << dimension + LOG(INFO) << "Test KVStateCache with tensorBytes: " << tensorBytes << ", capacity: " << capacity << ", layer: " << layer << ", block_size: " << block_size << "."; - init(dimension, capacity, layer, block_size, ipc_socket); + init(tensorBytes, capacity, layer, block_size, ipc_socket); for (size_t i = 0; i < tokens_list.size(); i++) { inference(tokens_list[i]); diff --git a/thirdparty/rax/radix.cc b/thirdparty/rax/radix.cc index 827396259..2a914b0b2 100644 --- a/thirdparty/rax/radix.cc +++ b/thirdparty/rax/radix.cc @@ -435,7 +435,7 @@ raxNode *raxAddChild(raxNode *n, int c, raxNode **childptr, raxNode ***parentlin * The function also returns a child node, since the last node of the * compressed chain cannot be part of the chain: it has zero children while * we can only compress inner nodes with exactly one child each. */ -raxNode *raxCompressNode(raxNode *n, int *s, size_t len, raxNode **child) { +raxNode *raxCompressNode(raxNode *n, const int *s, size_t len, raxNode **child) { assert(n->size == 0 && n->iscompr == 0); void *data = NULL; /* Initialized only to avoid warnings. */ size_t newsize; @@ -590,7 +590,7 @@ int handleOutOfMemory(rax *rax, raxNode *h, int *s, size_t len, void **old){ * function returns 0 as well but sets errno to ENOMEM, otherwise errno will * be set to 0. */ -int raxGenericInsert(rax *rax, int *s, size_t len, void *data, void **old, int overwrite, void **dataNode, bool set_timestamp = true) { +int raxGenericInsert(rax *rax, const int *s, size_t len, void *data, void **old, int overwrite, void **dataNode, bool set_timestamp = true) { size_t i; int j = 0; /* Split position. If raxLowWalk() stops in a compressed node, the index 'j' represents the char we stopped within the @@ -614,7 +614,6 @@ int raxGenericInsert(rax *rax, int *s, size_t len, void *data, void **old, int o debugf("### Insert: node representing key exists\n"); /* Make space for the value pointer if needed. */ if (!h->iskey || (h->isnull && overwrite)) { - printf("#############raxReallocForData1 ############\n"); h = raxReallocForData(h,data); if (h) memcpy(parentlink,&h,sizeof(h)); } @@ -1014,7 +1013,6 @@ int raxGenericInsert(rax *rax, int *s, size_t len, void *data, void **old, int o raxStackFree(&lowWalkStack); raxStackFree(&splitStack); raxNode *newh = raxReallocForData(h,data); - // printf("#############raxReallocForData2 ############\n"); if (newh == NULL) { return handleOutOfMemory(rax, h, (int *)s, i, old); } @@ -1028,7 +1026,7 @@ int raxGenericInsert(rax *rax, int *s, size_t len, void *data, void **old, int o /* Overwriting insert. Just a wrapper for raxGenericInsert() that will * update the element if there is already one for the same key. */ -int raxInsert(rax *rax, int *s, size_t len, void *data, void **old, bool set_timestamp) { +int raxInsert(rax *rax, const int *s, size_t len, void *data, void **old, bool set_timestamp) { void *dataNode = NULL; return raxGenericInsert(rax,s,len,data,old,1,&dataNode, set_timestamp); } @@ -1044,7 +1042,7 @@ int raxTryInsert(rax *rax, int *s, size_t len, void *data, void **old) { /* Overwriting insert. Return the raxNode that contains the key. */ -int raxInsertAndReturnDataNode(rax *rax, int *s, size_t len, void *data, void **node, void **old) { +int raxInsertAndReturnDataNode(rax *rax, const int *s, size_t len, void *data, void **node, void **old) { return raxGenericInsert(rax,s,len,data,old,1, node); } @@ -1065,7 +1063,7 @@ void *raxFind(rax *rax, int *s, size_t len) { /* ** Find a key in the rax, returns the stack */ -raxStack raxFindWithStack(rax *rax, int *s, size_t len) { +raxStack raxFindWithStack(rax *rax, const int *s, size_t len) { raxNode *h; raxStack ts; @@ -1081,12 +1079,11 @@ raxStack raxFindWithStack(rax *rax, int *s, size_t len) { /* ** Find a key in the rax, returns the raxNode that contains the key. */ -raxNode *raxFindAndReturnDataNode(rax *rax, int *s, size_t len, raxNode** sub_tree_node, bool set_timestamp) { +raxNode *raxFindAndReturnDataNode(rax *rax, const int *s, size_t len, raxNode** sub_tree_node, bool set_timestamp) { 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,set_timestamp); if (i != len || (h->iscompr && splitpos != 0) || !h->iskey) @@ -1102,25 +1099,6 @@ 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; @@ -1251,7 +1229,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, bool set_timestamp) { +int raxRemove(rax *rax, const int *s, size_t len, void **old, bool set_timestamp) { raxNode *h; raxStack ts; @@ -2169,33 +2147,36 @@ struct DebugTreeData { }; bool isPtr = true; }; + /* The actual implementation of raxShow(). */ -void raxRecursiveShow(int level, int lpad, raxNode *n) { +std::string raxRecursiveShow(int level, int lpad, raxNode *n) { + std::ostringstream raxStrStream; char s = n->iscompr ? '"' : '['; char e = n->iscompr ? '"' : ']'; - int numchars = printf("%c", s); + raxStrStream << s; for (int i = 0; i < n->size; i++) { - numchars += printf("%d ", n->data[i]); + raxStrStream << n->data[i] << " "; } - numchars += printf("%c %d ", e, n->numnodes); + + raxStrStream << e << " " << std::to_string(n->size) << " "; if (n->issubtree) { - numchars += printf("# "); - printf(" %p ", n); + raxStrStream << "# "; + raxStrStream << " " << n << " "; } if (n->iskey) { - numchars += printf("=%p",raxGetData(n)); + raxStrStream << "=" << raxGetData(n); } - numchars += printf(" node:%p time:%ld, data:%p, is_sub_tree:%d is_compr:%d", n, n->timestamp, n->custom_data, n->issubtree, n->iscompr); + raxStrStream << " node:" << n << " time:" << n->timestamp << " data:" << n->custom_data << " is_sub_tree:" << n->issubtree << " is_compr:" << n->iscompr; if (n->issubtree && n->custom_data != NULL) { - numchars += printf(" cus data:%p" , ((DebugDatawrapper *)(n->custom_data))->data); + raxStrStream << " cus data:" << ((DebugDatawrapper *)(n->custom_data))->data; DebugTreeData *data = (DebugTreeData *)((DebugDatawrapper *)(n->custom_data))->data; if (data) { if (data->isPtr) { - numchars += printf(" builder ptr:%p", data->kvStateCacheBlockBuilder); + raxStrStream << " builder ptr:" << data->kvStateCacheBlockBuilder; } else { - numchars += printf(" builder id:%lu", data->builderObjectID); + raxStrStream << " builder id:" << data->builderObjectID; } } } @@ -2205,31 +2186,33 @@ void raxRecursiveShow(int level, int lpad, raxNode *n) { * of " `-(x) " and " -> " respectively. */ if (level) { lpad += (numchildren > 1) ? 7 : 4; - if (numchildren == 1) lpad += numchars; + if (numchildren == 1) lpad += raxStrStream.str().size(); } raxNode **cp = raxNodeFirstChildPtr(n); for (int i = 0; i < numchildren; i++) { - const char *branch = " `-(%d) "; if (numchildren > 1) { - printf("\n"); - for (int j = 0; j < lpad; j++) putchar(' '); - printf(branch,n->data[i]); + raxStrStream << "\n"; + for (int j = 0; j < lpad; j++) raxStrStream << " "; + raxStrStream << " `-(" << n->data[i] << ") "; } else { - printf(" -> "); + raxStrStream << " -> "; } raxNode *child; memcpy(&child,cp,sizeof(child)); - raxRecursiveShow(level+1,lpad,child); + raxStrStream << raxRecursiveShow(level+1,lpad,child); cp++; } + return raxStrStream.str(); } /* 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'); +std::string raxShow(rax *rax) { + std::string raxStr; + raxStr += "rax numnode:" + std::to_string(rax->numele) + "\n"; + raxStr += raxRecursiveShow(0,0,rax->head); + raxStr += "\n"; + return raxStr; } /* Used by debugnode() macro to show info about a given node. */ @@ -2339,7 +2322,7 @@ bool raxIsSubtree(raxNode *node) { * tree from the root node. * */ -raxNode *raxSplit(rax *rax, int *s, size_t len, std::vector& token) { +raxNode *raxSplit(rax *rax, const int *s, size_t len, std::vector& token) { raxNode *childNode = NULL; raxNode *splitNode = NULL; raxStack stack = raxFindWithStack(rax, s, len); @@ -2490,12 +2473,12 @@ void raxFindLastRecentNode(raxNode *node, std::vector& key) { bool compareKey(int *first_key, int *second_key, int first_key_len, int second_key_len) { if (first_key_len != second_key_len) { - printf("length not equal, %d : %d\n", first_key_len, second_key_len); + VLOG(100) << "length not equal, " << first_key_len << " : " << second_key_len; return false; } for (int i = 0; i < first_key_len; i++) { if (first_key[i] != second_key[i]) { - printf("key not equal, %d : %d\n", first_key[i], second_key[i]); + VLOG(100) << "key not equal, " << first_key[i] << " : " << second_key[i]; return false; } } @@ -2513,12 +2496,16 @@ void sortNode(std::vector &ite_list) { std::sort(ite_list.begin(), ite_list.end(), compare); } -void printVector(int* v, int size) { - printf("token:\n"); - for (int i = 0; i < size; i++) { - printf("%d " , v[i]); +std::string printVector(std::vector& tree_iter_list) { + std::string token_str; + for (size_t i = 0; i < tree_iter_list.size(); i++) { + for (size_t j = 0; j < tree_iter_list[i].key_len; j++) { + token_str += std::to_string(tree_iter_list[i].key[j]); + token_str += " "; + } + token_str += "\n"; } - printf("\n"); + return token_str; } void freeVector(std::vector &ite_list) { @@ -2530,11 +2517,14 @@ void freeVector(std::vector &ite_list) { void mergeTree(rax* first_tree, rax* second_tree, std::vector>& evicted_tokens, std::set>& insert_tokens, int max_node) { - printf("merge tree!\n"); VLOG(100) << "==============tree 1===================="; - //raxShow(first_tree); + if (VLOG_IS_ON(100)) { + VLOG(100) << raxShow(first_tree); + } VLOG(100) << "==============tree 2===================="; - //raxShow(second_tree); + if (VLOG_IS_ON(100)) { + VLOG(100) << (second_tree); + } raxIterator first_tree_iter; raxIterator second_tree_iter; rax* tmp = raxNew(); @@ -2562,14 +2552,12 @@ void mergeTree(rax* first_tree, rax* second_tree, second_tree_iter_list.push_back(tmp_iter); } - for (size_t i = 0; i < first_tree_iter_list.size(); i++) { - printVector(first_tree_iter_list[i].key, first_tree_iter_list[i].key_len); - } - - for (size_t i = 0; i < second_tree_iter_list.size(); i++) { - printVector(second_tree_iter_list[i].key, second_tree_iter_list[i].key_len); + if (VLOG_IS_ON(100)) { + VLOG(100) << "==============first tree iter list====================" + << printVector(first_tree_iter_list); + VLOG(100) << "==============second tree iter list====================" + << printVector(second_tree_iter_list); } - // Sort by the length of the key, or timestamp if the keys have the same length. sortNode(first_tree_iter_list); sortNode(second_tree_iter_list); @@ -2587,7 +2575,7 @@ void mergeTree(rax* first_tree, rax* second_tree, second_tree_index == second_tree_iter_list.size()) { break; } - printf("nodeCount: %d\n", nodeCount); + VLOG(100) << "nodeCount: " << nodeCount; /** * If the key is the same, use the larger timestamp to refresh @@ -2598,8 +2586,7 @@ void mergeTree(rax* first_tree, rax* second_tree, second_tree_iter_list[second_tree_index].key, first_tree_iter_list[first_tree_index].key_len, second_tree_iter_list[second_tree_index].key_len)) { - // same key - printf("same key\n"); + VLOG(100) << "same key"; first_tree_iter_list[first_tree_index].node->timestamp = std::max(first_tree_iter_list[first_tree_index].node->timestamp, second_tree_iter_list[second_tree_index].node->timestamp); @@ -2624,9 +2611,9 @@ void mergeTree(rax* first_tree, rax* second_tree, * If the key is not in the record tree, it means that the key has not been * chosen in the past. So we need to insert the key into the record tree. */ - printf("choose first key %ld : %ld\n", - first_tree_iter_list[first_tree_index].node->timestamp, - second_tree_iter_list[second_tree_index].node->timestamp); + VLOG(100) << "choose first key, " + << "first tree iter timestamp: " << first_tree_iter_list[first_tree_index].node->timestamp + << "second tree iter timestamp: " << second_tree_iter_list[second_tree_index].node->timestamp; if (raxFind(tmp, first_tree_iter_list[first_tree_index].key, first_tree_iter_list[first_tree_index].key_len) == raxNotFound) { raxInsert(tmp, first_tree_iter_list[first_tree_index].key, @@ -2653,9 +2640,9 @@ void mergeTree(rax* first_tree, rax* second_tree, * chosen in the past. So we need to insert the key into the record tree. * and insert the key into the insert_tokens. */ - printf("choose second key %ld : %ld\n", - first_tree_iter_list[first_tree_index].node->timestamp, - second_tree_iter_list[second_tree_index].node->timestamp); + VLOG(100) << "choose second key, " + << "first tree iter timestamp: " << first_tree_iter_list[first_tree_index].node->timestamp + << "second tree iter timestamp: " << second_tree_iter_list[second_tree_index].node->timestamp; // choose second key if (raxFind(tmp, second_tree_iter_list[second_tree_index].key, second_tree_iter_list[second_tree_index].key_len) == raxNotFound) { @@ -2678,9 +2665,9 @@ void mergeTree(rax* first_tree, rax* second_tree, */ if (first_tree_iter_list[first_tree_index].node->numnodes <= second_tree_iter_list[second_tree_index].node->numnodes) { - printf("choose first key %ld : %ld\n", - first_tree_iter_list[first_tree_index].node->timestamp, - second_tree_iter_list[second_tree_index].node->timestamp); + VLOG(100) << "choose first key, " + << "first tree iter timestamp: " << first_tree_iter_list[first_tree_index].node->timestamp + << "second tree iter timestamp: " << second_tree_iter_list[second_tree_index].node->timestamp; // choose first key if (raxFind(tmp, first_tree_iter_list[first_tree_index].key, first_tree_iter_list[first_tree_index].key_len) == raxNotFound) { @@ -2702,9 +2689,9 @@ void mergeTree(rax* first_tree, rax* second_tree, } first_tree_index++; } else { - printf("choose second key %ld : %ld\n", - first_tree_iter_list[first_tree_index].node->timestamp, - second_tree_iter_list[second_tree_index].node->timestamp); + VLOG(100) << "choose second key, " + << "first tree iter timestamp: " << first_tree_iter_list[first_tree_index].node->timestamp + << "second tree iter timestamp: " << second_tree_iter_list[second_tree_index].node->timestamp; // choose second key if (raxFind(tmp, second_tree_iter_list[second_tree_index].key, second_tree_iter_list[second_tree_index].key_len) == raxNotFound) { @@ -2725,7 +2712,7 @@ void mergeTree(rax* first_tree, rax* second_tree, } if (nodeCount == max_node) { - printf("insert evicted tokens\n"); + VLOG(100) << "insert evicted tokens"; int evicted_node_count = 0; while (first_tree_index < first_tree_iter_list.size()) { if (raxFind(tmp, first_tree_iter_list[first_tree_index].key, @@ -2752,7 +2739,7 @@ void mergeTree(rax* first_tree, rax* second_tree, evicted_node_count++; first_tree_index++; } - printf("evicted_node_count: %d\n", evicted_node_count); + VLOG(100) << "evicted_node_count: " << evicted_node_count; freeVector(first_tree_iter_list); freeVector(second_tree_iter_list); raxFree(tmp); @@ -2790,9 +2777,8 @@ void mergeTree(rax* first_tree, rax* second_tree, } } else if (second_tree_index >= second_tree_iter_list.size()) { // second tree is empty - //raxShow(tmp); - printf("nodeCount:%d\n", nodeCount); - printf("first_tree_index:%ld\n", first_tree_index); + VLOG(100) << "nodeCount:" << nodeCount; + VLOG(100) << "first_tree_index:" << first_tree_index; while (first_tree_index < first_tree_iter_list.size() && nodeCount < max_node) { if (raxFind(tmp, first_tree_iter_list[first_tree_index].key, @@ -2844,26 +2830,6 @@ void mergeTree(rax* first_tree, rax* second_tree, raxFree(tmp); } -void testIteRax(rax *tree) { - raxIterator iter; - raxStart(&iter, tree); - raxSeek(&iter, "^", NULL, 0); - while (raxNext(&iter)) { - printf("key: "); - for (size_t i = 0; i < iter.key_len; i++) { - printf("%d ", iter.key[i]); - } - printf("\n"); - // printf("data: %p\n", iter.data); - } - raxStop(&iter); -} - raxNode* raxGetFirstChildPtr(raxNode* node) { return raxGetFirstChildPtr(node); } - -// 1 2 3 -// query subtree node:0x55f87076f760 -// I0129 16:44:25.626318 280948 kv_state_cache.cc:223] offset:0 -// I0129 16:44:25.626322 280948 kv_state_cache.cc:224] kvStateCacheBlockBuilder:0x55f870767bc0 \ No newline at end of file diff --git a/thirdparty/rax/radix.h b/thirdparty/rax/radix.h index 3590fb77f..63f8acdc1 100644 --- a/thirdparty/rax/radix.h +++ b/thirdparty/rax/radix.h @@ -210,15 +210,15 @@ extern void* raxNotFound; /* Exported API. */ rax* raxNew(void); -int raxInsert(rax* rax, int* s, size_t len, void* data, void** old, +int raxInsert(rax* rax, const int* s, size_t len, void* data, void** old, bool set_timestamp = true); int raxTryInsert(rax* rax, int* s, size_t len, void* data, void** old); -int raxInsertAndReturnDataNode(rax* rax, int* s, size_t len, void* data, +int raxInsertAndReturnDataNode(rax* rax, const int* s, size_t len, void* data, void** node, void** old); -int raxRemove(rax* rax, int* s, size_t len, void** old, +int raxRemove(rax* rax, const 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* raxFindAndReturnDataNode(rax* rax, const int* s, size_t len, raxNode** sub_tree_node = NULL, bool set_timestamp = true); void raxSetSubtree(raxNode* n); @@ -236,7 +236,7 @@ int raxRandomWalk(raxIterator* it, size_t steps); int raxCompare(raxIterator* iter, const char* op, int* key, size_t key_len); void raxStop(raxIterator* it); int raxEOF(raxIterator* it); -void raxShow(rax* rax); +std::string raxShow(rax* rax); uint64_t raxSize(rax* rax); void raxSetCustomData(raxNode* n, void* data); void* raxGetCustomData(raxNode* n); @@ -245,7 +245,7 @@ void raxSetDebugMsg(int onoff); void raxTraverse(raxNode* rax, std::vector>& dataNodeList); void raxTraverseSubTree(raxNode* n, std::vector& dataNodeList); -raxNode* raxSplit(rax* rax, int* s, size_t len, std::vector& key); +raxNode* raxSplit(rax* rax, const int* s, size_t len, std::vector& key); void raxSerialize(rax* root, std::vector>& tokenList, std::vector& dataList, std::vector& timestampsList, @@ -261,7 +261,5 @@ void raxFindLastRecentNode(raxNode* node, std::vector& key); void mergeTree(rax* first_tree, rax* second_tree, std::vector>& evicted_tokens, std::set>& insert_tokens, int max_node); -void testIteRax(rax* tree); raxNode* raxGetFirstChildPtr(raxNode* node); -// raxNode *raxSetSubTreeAndReturnDataNode(rax *rax, int *s, size_t len); #endif