Skip to content

Prefix caching for sequences with embeddings. #1841

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

Merged
ilya-lavrenov merged 9 commits into from
Mar 5, 2025
Merged

Prefix caching for sequences with embeddings. #1841

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
81f1263
Prefix caching for sequences with embeddings.
popovaan Mar 4, 2025
d053feb
Merge branch 'master' into prefix_caching_with_embeddings
popovaan Mar 4, 2025
e6ec494
Apply suggestions from code review
popovaan Mar 4, 2025
90b103c
Applied comments.
popovaan Mar 4, 2025
1d7dc19
Minor fix.
popovaan Mar 4, 2025
9b346ff
Minor correction.
popovaan Mar 4, 2025
f270c45
Minor correction.
popovaan Mar 4, 2025
fd7b11d
Minor correction.
popovaan Mar 4, 2025
a1a98d5
Merge branch 'master' into prefix_caching_with_embeddings
ilya-lavrenov Mar 4, 2025
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
15 changes: 7 additions & 8 deletions src/cpp/src/block_manager.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1073,7 +1073,7 @@ class BlockManager {
// When add_request() is executed in multiple threads accessing to cached_blocks causes segfault.
// The mutex is needed to prevent such segfaults.
const std::lock_guard<std::mutex> lock(m_cached_blocks_map_mutex);
auto prompt_ids = group->get_prompt_ids();
auto prompt_len = group->get_prompt_len();
auto sequences = group->get_not_finished_sequences();
OPENVINO_ASSERT(sequences.size() == 1);
auto sequence = sequences[0];
Expand All @@ -1085,11 +1085,11 @@ class BlockManager {
auto& block_table = m_block_table[seq_id];

size_t content_len = 0;
while (content_len < prompt_ids.size()) {
while (content_len < prompt_len) {
size_t prev_iteration_content_len = content_len;
content_len += m_block_size;
if (content_len > prompt_ids.size()) {
content_len = prompt_ids.size();
if (content_len > prompt_len) {
content_len = prompt_len;
}
// restore fully filled blocks
auto full_block_hash = sequence->get_hash(content_len);
Expand All @@ -1101,11 +1101,11 @@ class BlockManager {
block->set_timestamp(timestamp);
block_table[layer_idx].push_back(block);
}
group->update_processed_tokens_num(content_len == prompt_ids.size() ? content_len - 1 : content_len);
group->update_processed_tokens_num(content_len == prompt_len ? content_len - 1 : content_len);
} else {
// restore partially filled block
for (size_t i = 1; i < m_block_size; i++) {
if (prev_iteration_content_len + i > prompt_ids.size()) {
if (prev_iteration_content_len + i > prompt_len) {
break;
}
auto hash = sequence->get_hash(prev_iteration_content_len + i);
Expand All @@ -1118,8 +1118,7 @@ class BlockManager {
block->set_timestamp(timestamp);
block_table[layer_idx].push_back(block);
}

group->update_processed_tokens_num(prev_iteration_content_len + i == prompt_ids.size() ? prev_iteration_content_len + i - 1 : prev_iteration_content_len + i);
group->update_processed_tokens_num(prev_iteration_content_len + i == prompt_len ? prev_iteration_content_len + i - 1 : prev_iteration_content_len + i);

break;
}
Expand Down
7 changes: 4 additions & 3 deletions src/cpp/src/continuous_batching_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,9 +269,6 @@ ContinuousBatchingPipeline::ContinuousBatchingImpl::add_request(uint64_t request
SequenceGroup::Ptr sequence_group = std::make_shared<SequenceGroup>(request_id, input_ids, sampling_params, m_block_size);

if (m_scheduler->get_config().enable_prefix_caching) {
if (m_model_input_type == ModelInputType::EMBEDDINGS) {
OPENVINO_THROW("Prefix caching is not supported for VLM models.");
}
m_scheduler->restore_cached_blocks(sequence_group);
}

Expand Down Expand Up @@ -405,6 +402,10 @@ void ContinuousBatchingPipeline::ContinuousBatchingImpl::step() {

free_fork_timer.end();
}

// append embeddings for generated tokens
if (m_model_input_type == ModelInputType::EMBEDDINGS)
m_model_runner->append_embeddings(m_requests, scheduler_output);

// notify requests dropped by handle
{
Expand Down
87 changes: 59 additions & 28 deletions src/cpp/src/model_runner.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,6 @@ class ModelRunner {
size_t total_num_tokens = 0, total_num_blocks = 0;
size_t max_context_len_val = 0;
size_t hidden_size = 0;
size_t num_generated_ids = 0;
OPENVINO_ASSERT(sequence_groups.size() > 0);
auto sequence_group_type = sequence_groups[0]->get_sequence_group_type();
if (sequence_group_type == SequenceGroupType::EMBEDDINGS) {
Expand All @@ -135,9 +134,6 @@ class ModelRunner {
total_num_tokens += sequence_group->get_num_scheduled_tokens() * num_sequences;
total_num_blocks += sequence_group->get_num_blocks() * num_sequences;
max_context_len_val = std::max(max_context_len_val, sequence_group->get_context_len());
for (auto seq: sequence_group->get_running_sequences()) {
num_generated_ids += seq->get_generated_len();
}
}

ov::Tensor
Expand All @@ -163,27 +159,6 @@ class ModelRunner {
if (sequence_group_type == SequenceGroupType::EMBEDDINGS) {
OPENVINO_ASSERT(m_embedding.get_request(), "Got sequence group with embeddings, but embeddings model wasn't set.");
inputs_embeds_data = inputs_embeds.data<float>();

ov::Tensor generated_ids = ov::Tensor(ov::element::i64, {1, num_generated_ids});
int64_t *generated_ids_data = generated_ids.data<int64_t>();
size_t pos = 0;
for (size_t i = 0; i < num_sequence_groups; ++i) {
size_t seq_group_id = scheduler_output.m_scheduled_sequence_groups_ids[i];
SequenceGroup::CPtr sequence_group = sequence_groups[seq_group_id];
for (auto seq: sequence_group->get_running_sequences()) {
auto generated_ids = seq->get_generated_ids();
for (size_t token_idx = 0; token_idx < generated_ids.size(); token_idx++) {
generated_ids_data[pos] = generated_ids[token_idx];
pos++;
}
}
}
if (pos > 0) {
// TODO: Compute embeddings only for last generated token, while previously generated embeddings save in SequenceGroup
generated_ids_embeds = m_embedding.infer(generated_ids);
generated_ids_embeds_data = generated_ids_embeds.data<float>();
}

} else if (sequence_group_type == SequenceGroupType::TOKENS) {
input_ids_data = input_ids.data<int64_t>();
}
Expand Down Expand Up @@ -234,8 +209,8 @@ class ModelRunner {
sequence_group->get_prompt_ids()[position_id] :
sequence->get_generated_ids()[position_id - prompt_len];
} else if (sequence_group_type == SequenceGroupType::EMBEDDINGS) {
auto embeds_pos = position_id < prompt_len ? 0 : hidden_size * (position_id - prompt_len);
const float* src = position_id < prompt_len ? sequence_group->get_input_embeds()[position_id].data() : generated_ids_embeds_data + embeds_pos;
auto generated_embeds = sequence->get_generated_ids_embeds();
const float* src = position_id < prompt_len ? sequence_group->get_input_embeds()[position_id].data() : generated_embeds[position_id - prompt_len].data();
std::copy_n(src, hidden_size, inputs_embeds_data + token_id * hidden_size);
} else {
OPENVINO_THROW("Unknown model inputs type.");
Expand Down Expand Up @@ -271,7 +246,6 @@ class ModelRunner {
input_ids_data += num_scheduled_tokens;
} else if (sequence_group_type == SequenceGroupType::EMBEDDINGS) {
inputs_embeds_data += num_scheduled_tokens * hidden_size;
generated_ids_embeds_data += sequence->get_generated_len() * hidden_size;
}

position_ids_data += num_scheduled_tokens;
Expand Down Expand Up @@ -337,6 +311,63 @@ class ModelRunner {
return m_request.get_tensor("logits");
}

void append_embeddings(const std::vector<SequenceGroup::Ptr> & sequence_groups, const Scheduler::Output& scheduler_output) {
size_t num_sequence_groups = scheduler_output.m_scheduled_sequence_groups_ids.size();
size_t num_generated_ids_without_embeddings = 0;
OPENVINO_ASSERT(sequence_groups.size() > 0);

// compute aggregated values
for (size_t i = 0; i < num_sequence_groups; ++i) {
size_t seq_group_id = scheduler_output.m_scheduled_sequence_groups_ids[i];
SequenceGroup::CPtr sequence_group = sequence_groups[seq_group_id];
size_t num_sequences = sequence_group->num_running_seqs();
OPENVINO_ASSERT(sequence_group->get_sequence_group_type() == SequenceGroupType::EMBEDDINGS);
for (auto seq: sequence_group->get_running_sequences()) {
num_generated_ids_without_embeddings += seq->get_generated_len() - seq->get_generated_ids_embeds().size();
}
}
size_t hidden_size = sequence_groups[0]->get_hidden_size();

ov::Tensor generated_ids_embeds;
float *generated_ids_embeds_data = nullptr;

OPENVINO_ASSERT(m_embedding.get_request(), "Got sequence group with embeddings, but embeddings model wasn't set.");

ov::Tensor generated_ids = ov::Tensor(ov::element::i64, {1, num_generated_ids_without_embeddings});
int64_t *generated_ids_data = generated_ids.data<int64_t>();
size_t pos = 0;
for (size_t i = 0; i < num_sequence_groups; ++i) {
size_t seq_group_id = scheduler_output.m_scheduled_sequence_groups_ids[i];
SequenceGroup::CPtr sequence_group = sequence_groups[seq_group_id];
for (auto seq: sequence_group->get_running_sequences()) {
auto generated_ids = seq->get_generated_ids();
for (size_t token_idx = seq->get_generated_ids_embeds().size(); token_idx < generated_ids.size(); token_idx++) {
generated_ids_data[pos] = generated_ids[token_idx];
pos++;
}
}
}
if (pos > 0) {
generated_ids_embeds = m_embedding.infer(generated_ids);
generated_ids_embeds_data = generated_ids_embeds.data<float>();

for (size_t i = 0; i < num_sequence_groups; ++i) {
size_t seq_group_id = scheduler_output.m_scheduled_sequence_groups_ids[i];
size_t embeds_pos = 0;
SequenceGroup::Ptr sequence_group = sequence_groups[seq_group_id];
for (auto seq: sequence_group->get_running_sequences()) {
auto generated_ids = seq->get_generated_ids();
size_t new_embeds_count = seq->get_generated_len() - seq->get_generated_ids_embeds().size();
ov::Coordinate start{0, embeds_pos, 0};
ov::Coordinate end{1, embeds_pos + new_embeds_count, hidden_size};
ov::Tensor embedding(generated_ids_embeds, start, end);
seq->append_generated_ids_embeds(embedding);
embeds_pos += new_embeds_count;
}
}
}
}

private:
void _fill_indices_from_block_tables(
const std::vector<std::string>& dst_tensor_names,
Expand Down
74 changes: 65 additions & 9 deletions src/cpp/src/sequence_group.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,22 +22,78 @@ size_t Sequence::_make_hash(size_t content_length) {
size_t prefix_hashes_needed_count = block_start_idx / block_size;
OPENVINO_ASSERT(prefix_hashes_needed_count <= m_prefix_hashes.size());
content.insert(content.end(), m_prefix_hashes.begin(), m_prefix_hashes.begin() + prefix_hashes_needed_count);
char* data;
std::size_t size;

// get tokens corresponding to current block
const auto prompt_ids = sequence_group->get_prompt_ids();
OPENVINO_ASSERT(content_length <= prompt_ids.size() + m_generated_ids.size());
if (block_start_idx < prompt_ids.size()) {
content.insert(content.end(), prompt_ids.begin() + block_start_idx, prompt_ids.begin() + std::min(prompt_ids.size(), content_length));
if (sequence_group->get_sequence_group_type() == SequenceGroupType::TOKENS) {
const auto prompt_ids = sequence_group->get_prompt_ids();
OPENVINO_ASSERT(content_length <= prompt_ids.size() + m_generated_ids.size());
if (block_start_idx < prompt_ids.size()) {
content.insert(content.end(), prompt_ids.begin() + block_start_idx, prompt_ids.begin() + std::min(prompt_ids.size(), content_length));
}
if (content_length > prompt_ids.size()) {
size_t start = block_start_idx < prompt_ids.size() ? 0 : block_start_idx - prompt_ids.size();
content.insert(content.end(), m_generated_ids.begin() + start, m_generated_ids.begin() + content_length - prompt_ids.size());
}
data = reinterpret_cast<char*>(content.data());
size = content.size() * sizeof(content[0]);
}
if (content_length > prompt_ids.size()) {
size_t start = block_start_idx < prompt_ids.size() ? 0 : block_start_idx - prompt_ids.size();
content.insert(content.end(), m_generated_ids.begin() + start, m_generated_ids.begin() + content_length - prompt_ids.size());
else if (sequence_group->get_sequence_group_type() == SequenceGroupType::EMBEDDINGS) {
const auto input_embeds = sequence_group->get_input_embeds();
const auto generated_embeds = m_generated_ids_embeds;
OPENVINO_ASSERT(content_length <= input_embeds.size() + generated_embeds.size());
std::vector<float> content_float;

// get inputs embeddings
if (block_start_idx < input_embeds.size()) {
for (size_t idx = block_start_idx; idx < std::min(input_embeds.size(), content_length); idx++) {
auto embed = _reduce_embedding(input_embeds[idx]);
const char* embed_char = reinterpret_cast<const char*>(embed.data());
content_float.insert(content_float.end(), embed.begin(), embed.end());
}
}

// get generated ids embeddings
if (content_length > input_embeds.size()) {
size_t start = block_start_idx < input_embeds.size() ? 0 : block_start_idx - input_embeds.size();
for (size_t idx = start; idx < content_length - input_embeds.size(); idx++) {
auto embed = _reduce_embedding(generated_embeds[idx]);
content_float.insert(content_float.end(), embed.begin(), embed.end());
}
}

size_t prev_hashes_size = content.size() == 0 ? 0 : content.size() * sizeof(content[0]);
size_t content_float_size = content_float.size() * sizeof(content_float[0]);
size = prev_hashes_size + content_float_size;
data = new char[size];

// append previously calculated prefix hashes if they are available
if (prev_hashes_size) {
auto prev_hashes = reinterpret_cast<const char*>(content.data());
std::copy_n(prev_hashes, prev_hashes_size, data);
}

auto content_char = reinterpret_cast<const char*>(content_float.data());
std::copy_n(content_char, content_float_size, data + prev_hashes_size);
}
const char* data = reinterpret_cast<const char*>(content.data());
std::size_t size = content.size() * sizeof(content[0]);
else {
OPENVINO_THROW("Hash calculation is not supported for this sequence type.");
}
auto hash = std::hash<std::string_view>{}(std::string_view(data, size));
return std::hash<std::string_view>{}(std::string_view(data, size));
}

std::vector<float> Sequence::_reduce_embedding(const std::vector<float>& embedding) {
size_t s = embedding.size();
size_t res_size = std::min((size_t)ceil(float(embedding.size()) / m_embeddings_hash_calculation_stride), m_embeddings_hash_max_num_values);
std::vector<float> res(res_size);
for (size_t i = 0, idx=0; idx < res_size; i+= m_embeddings_hash_calculation_stride, idx++) {
res[idx] = embedding[i];
}
return res;
}

// Each KV block can be uniquely identified by
// the tokens within the block and the tokens in the prefix before the block.
// hash(prefix tokens + block tokens) <--> KV Block
Expand Down
Loading
Loading