Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support layer for kv state cache #1766

Merged
merged 3 commits into from
Feb 22, 2024
Merged

Support layer for kv state cache #1766

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7e3a88a
Support layer for kv state cache.
vegetableysm Feb 22, 2024
1e84581
Format code.
vegetableysm Feb 22, 2024
f446e46
Add multi layer test and fix bug.
vegetableysm Feb 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 9 additions & 15 deletions modules/kv-state-cache/ds/kv_state_cache.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ void KVStateCache::Resolve() {
// 3. construct the member field
this->dimension = this->meta_.GetKeyValue<int>("dimension");
this->version = this->meta_.GetKeyValue<uint64_t>("version");
this->layer = this->meta_.GetKeyValue<int>("layer");
LOG(INFO) << "construct the member field success" << std::endl;
}

Expand All @@ -68,11 +69,12 @@ KVStateCache::~KVStateCache() {
}

KVStateCacheBuilder::KVStateCacheBuilder(Client& client, int dimension,
int cacheCapacity) {
int cacheCapacity, int layer) {
this->dimension = dimension;
this->version = 0;
this->layer = layer;
KVStateCacheBlockBuilder* builder =
new KVStateCacheBlockBuilder(client, this->dimension);
new KVStateCacheBlockBuilder(client, this->dimension, layer);

this->rootTree = std::make_shared<RadixTree>(cacheCapacity);

Expand All @@ -95,6 +97,7 @@ KVStateCacheBuilder::KVStateCacheBuilder(Client& client,
// TBD
this->dimension = cache->GetDemension();
this->version = cache->GetVersion();
this->layer = cache->GetLayer();
// 1. create block builder from block
std::map<uint64_t, std::shared_ptr<KVStateCacheBlock>> kvStateCacheBlockMap =
cache->kvStateCacheBlockMap;
Expand Down Expand Up @@ -123,26 +126,16 @@ KVStateCacheBlockBuilder* KVStateCacheBuilder::Split(
// Split the tree if the list of kvState is full.
VINEYARD_ASSERT(nodeDataList.size() > 0);
KVStateCacheBlockBuilder* childKVStateCacheBlockBuilder =
new KVStateCacheBlockBuilder(client, this->dimension);
new KVStateCacheBlockBuilder(client, this->dimension, this->layer);
for (size_t i = 0; i < nodeDataList.size(); i++) {
OffsetData* data = (OffsetData*) nodeDataList[i]->nodeData->data;
if (data == nullptr)
continue;
int index = data->offset;

// Transfer the data from this builder to the child builder.
const std::shared_ptr<TensorBuilder<double>> keyStateTensorBuilder =
kvStateCacheBlockBuilder->GetKeyStateBuilder();
const std::shared_ptr<TensorBuilder<double>> valueStateTensorBuilder =
kvStateCacheBlockBuilder->GetValueStateBuilder();
OffsetData new_offset_data;
childKVStateCacheBlockBuilder->Update(
keyStateTensorBuilder->data() + index * this->dimension,
valueStateTensorBuilder->data() + index * this->dimension,
this->dimension, &new_offset_data);
data->offset = new_offset_data.offset;
// Clear the bitmap.
kvStateCacheBlockBuilder->DeleteKVCache(index);
data->offset =
kvStateCacheBlockBuilder->Split(childKVStateCacheBlockBuilder, index);
}
LOG(INFO) << "builder:" << kvStateCacheBlockBuilder
<< " bitmap:" << kvStateCacheBlockBuilder->GetBitmapStr();
Expand Down Expand Up @@ -330,6 +323,7 @@ std::shared_ptr<Object> KVStateCacheBuilder::_Seal(Client& client) {
// 1. store the member variables to cache object meta
kvStateCache->meta_.AddKeyValue("dimension", this->dimension);
kvStateCache->meta_.AddKeyValue("version", this->version);
kvStateCache->meta_.AddKeyValue("layer", this->layer);

// 2. seal all the block and put object id to cache object and
// change the tree data from pointer to object id
Expand Down
9 changes: 8 additions & 1 deletion modules/kv-state-cache/ds/kv_state_cache.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ class KVStateCache : public vineyard::Registered<KVStateCache> {
std::shared_ptr<RadixTree> rootTree;
int dimension;
int cacheCapacity;
int layer;
uint64_t version;

public:
Expand All @@ -68,6 +69,8 @@ class KVStateCache : public vineyard::Registered<KVStateCache> {

std::shared_ptr<RadixTree> GetRootTree() { return this->rootTree; }

int GetLayer() { return this->layer; }

~KVStateCache();

friend class KVStateCacheBuilder;
Expand All @@ -76,10 +79,12 @@ class KVStateCache : public vineyard::Registered<KVStateCache> {
class KVStateCacheBuilder : public vineyard::ObjectBuilder {
std::shared_ptr<RadixTree> rootTree;
int dimension;
int layer = 1;
uint64_t version;

public:
KVStateCacheBuilder(Client& client, int dimension, int cacheCapacity);
KVStateCacheBuilder(Client& client, int dimension, int cacheCapacity,
int layer);

KVStateCacheBuilder(Client& client, std::shared_ptr<KVStateCache> cache);

Expand Down Expand Up @@ -109,6 +114,8 @@ class KVStateCacheBuilder : public vineyard::ObjectBuilder {

std::shared_ptr<RadixTree> GetRootTree() { return this->rootTree; }

int GetLayer() { return this->layer; }

~KVStateCacheBuilder();
};

Expand Down
173 changes: 107 additions & 66 deletions modules/kv-state-cache/ds/kv_state_cache_block.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,63 +49,79 @@ void KVStateCacheBlock::Construct(const ObjectMeta& meta) {

// TBD
// 1. construct the keyStateTensorBuilder and valueStateTensorBuilder
this->keyStateTensor = std::dynamic_pointer_cast<Tensor<double>>(
this->meta_.GetMember("keyStateTensorBuilder"));
this->valueStateTensor = std::dynamic_pointer_cast<Tensor<double>>(
this->meta_.GetMember("valueStateTensorBuilder"));
this->layer = this->meta_.GetKeyValue<int>("layer");
for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) {
this->keyStateTensorList.push_back(
std::dynamic_pointer_cast<Tensor<double>>(this->meta_.GetMember(
"keyStateTensorBuilder_" + std::to_string(currentLayer))));
this->valueStateTensorList.push_back(
std::dynamic_pointer_cast<Tensor<double>>(this->meta_.GetMember(
"valueStateTensorBuilder_" + std::to_string(currentLayer))));
}
// 2. construct the member field
this->bitmap = this->meta_.GetKeyValue<unsigned long long>("bitmap");
this->dimension = this->meta_.GetKeyValue<int>("dimension");
}

KVStateCacheBlockBuilder::KVStateCacheBlockBuilder(Client& client,
int dimension) {
int dimension, int layer) {
this->bitmap = UINT64_MAX;
std::vector<int64_t> shape = {LIST_SIZE, dimension};
this->keyStateTensorBuilder =
std::make_shared<TensorBuilder<double>>(client, shape);
this->valueStateTensorBuilder =
std::make_shared<TensorBuilder<double>>(client, shape);
for (int i = 0; i < layer; i++) {
this->keyStateTensorBuilderList.push_back(
std::make_shared<TensorBuilder<double>>(client, shape));
this->valueStateTensorBuilderList.push_back(
std::make_shared<TensorBuilder<double>>(client, shape));
}
this->dimension = dimension;
this->layer = layer;
}

KVStateCacheBlockBuilder::KVStateCacheBlockBuilder(
Client& client, std::shared_ptr<KVStateCacheBlock> kvStateCacheBlock) {
this->bitmap = kvStateCacheBlock->bitmap;
this->dimension = kvStateCacheBlock->dimension;
this->layer = kvStateCacheBlock->layer;
std::vector<int64_t> shape = {LIST_SIZE, dimension};
this->keyStateTensorBuilder =
std::make_shared<TensorBuilder<double>>(client, shape);
this->valueStateTensorBuilder =
std::make_shared<TensorBuilder<double>>(client, shape);

// transfer the data from kv_state_cache to this builder
memcpy(this->keyStateTensorBuilder->data(),
kvStateCacheBlock->keyStateTensor->data(),
LIST_SIZE * this->dimension * sizeof(double));
memcpy(this->valueStateTensorBuilder->data(),
kvStateCacheBlock->valueStateTensor->data(),
LIST_SIZE * this->dimension * sizeof(double));
for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) {
this->keyStateTensorBuilderList.push_back(
std::make_shared<TensorBuilder<double>>(client, shape));
this->valueStateTensorBuilderList.push_back(
std::make_shared<TensorBuilder<double>>(client, shape));
}

for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) {
memcpy(this->keyStateTensorBuilderList[currentLayer]->data(),
kvStateCacheBlock->keyStateTensorList[currentLayer]->data(),
LIST_SIZE * this->dimension * sizeof(double));
memcpy(this->valueStateTensorBuilderList[currentLayer]->data(),
kvStateCacheBlock->valueStateTensorList[currentLayer]->data(),
LIST_SIZE * this->dimension * sizeof(double));
}
}

// current we do not consider the layer.
Status KVStateCacheBlockBuilder::Query(Client& client, int index,
KV_STATE_WITH_LAYER& kvState) {
std::vector<double> keyStateVector;
std::vector<double> valueStateVector;

for (int i = 0; i < this->dimension; ++i) {
keyStateVector.push_back(
((double*) keyStateTensorBuilder->data())[index * dimension + i]);
}

for (int i = 0; i < this->dimension; ++i) {
valueStateVector.push_back(
((double*) valueStateTensorBuilder->data())[index * dimension + i]);
for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) {
std::vector<double> keyStateVector;
std::vector<double> valueStateVector;

for (int i = 0; i < this->dimension; ++i) {
keyStateVector.push_back(
((double*) keyStateTensorBuilderList[currentLayer]
->data())[index * dimension + i]);
}

for (int i = 0; i < this->dimension; ++i) {
valueStateVector.push_back(
((double*) valueStateTensorBuilderList[currentLayer]
->data())[index * dimension + i]);
}

kvState.insert(std::make_pair(
currentLayer, std::make_pair(keyStateVector, valueStateVector)));
}

kvState.insert(
std::make_pair(1, std::make_pair(keyStateVector, valueStateVector)));
return Status::OK();
}

Expand All @@ -124,40 +140,60 @@ void KVStateCacheBlockBuilder::Update(const KV_STATE_WITH_LAYER& kvState,
OffsetData* data) {
int index = this->FindEmptySlot();
LOG(INFO) << "index:" << index;
std::vector<double> keyStateVector = (kvState.find(1)->second).first;
std::vector<double> valueStateVector = (kvState.find(1)->second).second;
VINEYARD_ASSERT(keyStateVector.size() == (size_t) this->dimension);
VINEYARD_ASSERT(valueStateVector.size() == (size_t) this->dimension);

double* keyData = (double*) keyStateTensorBuilder->data();
double* valueData = (double*) valueStateTensorBuilder->data();
for (int i = 0; i < this->dimension; ++i) {
keyData[index * this->dimension + i] = keyStateVector[i];
}
for (int i = 0; i < this->dimension; ++i) {
valueData[index * this->dimension + i] = valueStateVector[i];
LOG(INFO) << "layer:" << layer;
for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) {
std::vector<double> keyStateVector =
(kvState.find(currentLayer)->second).first;
std::vector<double> valueStateVector =
(kvState.find(currentLayer)->second).second;
LOG(INFO) << "vector size:" << keyStateVector.size() << " "
<< valueStateVector.size() << " demension" << this->dimension;
VINEYARD_ASSERT(keyStateVector.size() == (size_t) this->dimension);
VINEYARD_ASSERT(valueStateVector.size() == (size_t) this->dimension);

double* keyData = (double*) keyStateTensorBuilderList[currentLayer]->data();
double* valueData =
(double*) valueStateTensorBuilderList[currentLayer]->data();
memcpy(keyData + index * this->dimension, keyStateVector.data(),
this->dimension * sizeof(double));
memcpy(valueData + index * this->dimension, valueStateVector.data(),
this->dimension * sizeof(double));
}
data->offset = index;

ACQUIRE_BIT_RESOURCE(this->bitmap, index);
}

void KVStateCacheBlockBuilder::Update(double* keyState, double* valueState,
unsigned long dataLength,
OffsetData* data) {
int index = FindEmptySlot();
double* keyData = (double*) keyStateTensorBuilder->data();
double* valueData = (double*) valueStateTensorBuilder->data();
VINEYARD_ASSERT((unsigned long) this->dimension == dataLength);
for (unsigned long i = 0; i < dataLength; ++i) {
keyData[index * this->dimension + i] = keyState[i];
}
for (unsigned long i = 0; i < dataLength; ++i) {
valueData[index * this->dimension + i] = valueState[i];
short KVStateCacheBlockBuilder::Split(KVStateCacheBlockBuilder* child,
int index) {
// TBD
VINEYARD_ASSERT(this->layer == child->layer);
int childIndex = child->FindEmptySlot();
for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) {
std::shared_ptr<TensorBuilder<double>> keyStateTensorBuilder =
keyStateTensorBuilderList[currentLayer];
std::shared_ptr<TensorBuilder<double>> valueStateTensorBuilder =
valueStateTensorBuilderList[currentLayer];
std::shared_ptr<TensorBuilder<double>> childKeyStateTensorBuilder =
child->keyStateTensorBuilderList[currentLayer];
std::shared_ptr<TensorBuilder<double>> childValueStateTensorBuilder =
child->valueStateTensorBuilderList[currentLayer];

double* keyState =
(double*) keyStateTensorBuilder->data() + index * this->dimension;
double* valueState =
(double*) valueStateTensorBuilder->data() + index * this->dimension;
double* childKeyState = (double*) childKeyStateTensorBuilder->data() +
childIndex * this->dimension;
double* childValueState = (double*) childValueStateTensorBuilder->data() +
childIndex * this->dimension;

memcpy(childKeyState, keyState, this->dimension * sizeof(double));
memcpy(childValueState, valueState, this->dimension * sizeof(double));
}
data->offset = index;

ACQUIRE_BIT_RESOURCE(this->bitmap, index);
ACQUIRE_BIT_RESOURCE(child->bitmap, childIndex);
FREE_BIT_RESOURCE(this->bitmap, index);
return childIndex;
}

Status KVStateCacheBlockBuilder::Build(Client& client) { return Status::OK(); }
Expand All @@ -170,14 +206,19 @@ std::shared_ptr<Object> KVStateCacheBlockBuilder::_Seal(Client& client) {
std::make_shared<KVStateCacheBlock>();

// 1. seal keyStateTensorBuilder and valueStateTensorBuilder
kvStateCacheBlock->meta_.AddMember("keyStateTensorBuilder",
keyStateTensorBuilder->Seal(client));
kvStateCacheBlock->meta_.AddMember("valueStateTensorBuilder",
valueStateTensorBuilder->Seal(client));
for (int currentLayer = 0; currentLayer < this->layer; currentLayer++) {
kvStateCacheBlock->meta_.AddMember(
"keyStateTensorBuilder_" + std::to_string(currentLayer),
keyStateTensorBuilderList[currentLayer]->Seal(client));
kvStateCacheBlock->meta_.AddMember(
"valueStateTensorBuilder_" + std::to_string(currentLayer),
valueStateTensorBuilderList[currentLayer]->Seal(client));
}

// 2. store the member field to meta
kvStateCacheBlock->meta_.AddKeyValue("bitmap", this->bitmap);
kvStateCacheBlock->meta_.AddKeyValue("dimension", this->dimension);
kvStateCacheBlock->meta_.AddKeyValue("layer", this->layer);
// 3. set the object type to meta
kvStateCacheBlock->meta_.SetTypeName(type_name<KVStateCacheBlock>());

Expand Down
Loading
Loading