diff --git a/src/cpp/src/cache_manager.hpp b/src/cpp/src/cache_manager.hpp index 0c04823f4f..20d4c0c51c 100644 --- a/src/cpp/src/cache_manager.hpp +++ b/src/cpp/src/cache_manager.hpp @@ -46,8 +46,6 @@ class CacheManager { } OPENVINO_ASSERT(m_key_cache.size() == m_value_cache.size()); m_num_allocated_kv_blocks = num_kv_blocks; - ov::Shape value_cache_shape = set_first_dim_and_make_static(m_device_config.get_value_cache_shape(), num_kv_blocks); - ov::Shape key_cache_shape = set_first_dim_and_make_static(m_device_config.get_key_cache_shape(), num_kv_blocks); const std::string device_name = m_device_config.get_device(); @@ -56,6 +54,8 @@ class CacheManager { if (device_name.find("GPU") == std::string::npos) {// Allocate KV caches for (size_t decoder_layer_id = 0; decoder_layer_id < m_device_config.get_num_layers(); ++decoder_layer_id) { + ov::Shape value_cache_shape = set_first_dim_and_make_static(m_device_config.get_value_cache_shape(decoder_layer_id), num_kv_blocks); + ov::Shape key_cache_shape = set_first_dim_and_make_static(m_device_config.get_key_cache_shape(decoder_layer_id), num_kv_blocks); ov::Tensor key_cache(m_device_config.get_cache_precision(), key_cache_shape); ov::Tensor value_cache(m_device_config.get_cache_precision(), value_cache_shape); @@ -104,6 +104,8 @@ class CacheManager { } else { auto remote_context = m_core.get_default_context(device_name); for (size_t decoder_layer_id = 0; decoder_layer_id < m_device_config.get_num_layers(); ++decoder_layer_id) { + ov::Shape value_cache_shape = set_first_dim_and_make_static(m_device_config.get_value_cache_shape(decoder_layer_id), num_kv_blocks); + ov::Shape key_cache_shape = set_first_dim_and_make_static(m_device_config.get_key_cache_shape(decoder_layer_id), num_kv_blocks); ov::Tensor key_cache = remote_context.create_tensor(m_device_config.get_cache_precision(), key_cache_shape); ov::Tensor value_cache = remote_context.create_tensor(m_device_config.get_cache_precision(), @@ -142,30 +144,27 @@ class CacheManager { } void copy_blocks(const std::map>& block_copy_map) { - ov::Shape key_shape = set_first_dim_and_make_static(m_device_config.get_key_cache_shape(), m_num_allocated_kv_blocks); - ov::Shape value_shape = set_first_dim_and_make_static(m_device_config.get_value_cache_shape(), m_num_allocated_kv_blocks); - - ov::Coordinate key_src_start_roi(key_shape.size(), 0); - ov::Coordinate key_src_end_roi = key_shape; - ov::Coordinate key_dst_start_roi(key_shape.size(), 0); - ov::Coordinate key_dst_end_roi = key_shape; - - ov::Coordinate value_src_start_roi(value_shape.size(), 0); - ov::Coordinate value_src_end_roi = value_shape; - ov::Coordinate value_dst_start_roi(value_shape.size(), 0); - ov::Coordinate value_dst_end_roi = value_shape; - for (const auto & blocks_pair : block_copy_map) { size_t src_block_id = blocks_pair.first; - key_src_end_roi[0] = (key_src_start_roi[0] = src_block_id) + 1; - value_src_end_roi[0] = (value_src_start_roi[0] = src_block_id) + 1; - const std::list& dst_block_ids = blocks_pair.second; for (size_t dst_block_id : dst_block_ids) { - key_dst_end_roi[0] = (key_dst_start_roi[0] = dst_block_id) + 1; - value_dst_end_roi[0] = (value_dst_start_roi[0] = dst_block_id) + 1; - for (size_t decoder_layer_id = 0; decoder_layer_id < m_device_config.get_num_layers(); ++decoder_layer_id) { + ov::Shape key_shape = set_first_dim_and_make_static(m_device_config.get_key_cache_shape(decoder_layer_id), m_num_allocated_kv_blocks); + ov::Shape value_shape = set_first_dim_and_make_static(m_device_config.get_value_cache_shape(decoder_layer_id), m_num_allocated_kv_blocks); + ov::Coordinate key_src_start_roi(key_shape.size(), 0); + ov::Coordinate key_src_end_roi = key_shape; + ov::Coordinate key_dst_start_roi(key_shape.size(), 0); + ov::Coordinate key_dst_end_roi = key_shape; + + ov::Coordinate value_src_start_roi(value_shape.size(), 0); + ov::Coordinate value_src_end_roi = value_shape; + ov::Coordinate value_dst_start_roi(value_shape.size(), 0); + ov::Coordinate value_dst_end_roi = value_shape; + key_src_end_roi[0] = (key_src_start_roi[0] = src_block_id) + 1; + value_src_end_roi[0] = (value_src_start_roi[0] = src_block_id) + 1; + key_dst_end_roi[0] = (key_dst_start_roi[0] = dst_block_id) + 1; + value_dst_end_roi[0] = (value_dst_start_roi[0] = dst_block_id) + 1; + ov::Tensor key_src_cache_roi(m_key_cache[decoder_layer_id], key_src_start_roi, key_src_end_roi); ov::Tensor key_dst_cache_roi(m_key_cache[decoder_layer_id], key_dst_start_roi, key_dst_end_roi); diff --git a/src/cpp/src/device_config.hpp b/src/cpp/src/device_config.hpp index 371142701c..cc2e21b9a1 100644 --- a/src/cpp/src/device_config.hpp +++ b/src/cpp/src/device_config.hpp @@ -12,8 +12,9 @@ namespace ov::genai { class DeviceConfig { ov::element::Type m_kv_cache_type; - ov::PartialShape m_key_cache_shape, m_value_cache_shape; - ov::Shape::value_type m_num_kv_heads, m_head_size, m_num_decoder_layers; + std::vector m_key_cache_shape, m_value_cache_shape; + std::vector m_num_kv_heads; + ov::Shape::value_type m_head_size, m_num_decoder_layers; size_t m_num_kv_blocks = 0; size_t m_block_size = 0; size_t m_cache_size = 0; @@ -88,11 +89,14 @@ class DeviceConfig { } } - void set_model_params(size_t num_kv_heads, size_t head_size, size_t num_decoder_layers) { - m_num_kv_heads = num_kv_heads; + void set_model_params(std::vector num_kv_heads, size_t head_size, size_t num_decoder_layers) { m_head_size = head_size; m_num_decoder_layers = num_decoder_layers; + m_num_kv_heads.assign(num_kv_heads.begin(), num_kv_heads.end()); + m_key_cache_shape.reserve(m_num_decoder_layers); + m_value_cache_shape.reserve(m_num_decoder_layers); + if (m_device == "CPU") { // Scale, zero point and quantized data will be stored together. // The layout for per token per head: @@ -104,21 +108,32 @@ class DeviceConfig { } if (m_num_kv_blocks == 0 && m_cache_size > 0) { + size_t block_size = 0; size_t size_in_bytes = m_cache_size * 1024 * 1024 * 1024; - m_num_kv_blocks = size_in_bytes / (m_num_decoder_layers * 2 * m_num_kv_heads * m_block_size * m_head_size * m_kv_cache_type.size()); + for (size_t layer_id = 0; layer_id < m_num_decoder_layers; layer_id++) { + block_size += 2 * m_num_kv_heads[layer_id] * m_block_size * m_head_size * m_kv_cache_type.size(); + } + m_num_kv_blocks = size_in_bytes / block_size; } - m_key_cache_shape = m_value_cache_shape = ov::PartialShape{ov::Dimension::dynamic(), - ov::Dimension(m_num_kv_heads), - ov::Dimension(m_block_size), - ov::Dimension(m_head_size)}; - - if (m_device.find("GPU") != std::string::npos) { - // Update key shape, as the key's shape is different from the value's shape - m_key_cache_shape = ov::PartialShape{ov::Dimension::dynamic(), - ov::Dimension(m_num_kv_heads), - ov::Dimension(m_head_size), - ov::Dimension(m_block_size)}; + for (size_t layer_id = 0; layer_id < m_num_decoder_layers; layer_id++) { + m_key_cache_shape.push_back(ov::PartialShape{ov::Dimension::dynamic(), + ov::Dimension(m_num_kv_heads[layer_id]), + ov::Dimension(m_block_size), + ov::Dimension(m_head_size)}); + + m_value_cache_shape.push_back(ov::PartialShape{ov::Dimension::dynamic(), + ov::Dimension(m_num_kv_heads[layer_id]), + ov::Dimension(m_block_size), + ov::Dimension(m_head_size)}); + + if (m_device.find("GPU") != std::string::npos) { + // Update key shape, as the key's shape is different from the value's shape + m_key_cache_shape.push_back(ov::PartialShape{ov::Dimension::dynamic(), + ov::Dimension(m_num_kv_heads[layer_id]), + ov::Dimension(m_head_size), + ov::Dimension(m_block_size)}); + } } } @@ -134,14 +149,14 @@ class DeviceConfig { return m_num_decoder_layers; } - ov::PartialShape get_key_cache_shape() const { + ov::PartialShape get_key_cache_shape(size_t id) const { OPENVINO_ASSERT(m_key_cache_shape.size()); - return m_key_cache_shape; + return m_key_cache_shape[id]; } - ov::PartialShape get_value_cache_shape() const { + ov::PartialShape get_value_cache_shape(size_t id) const { OPENVINO_ASSERT(m_value_cache_shape.size()); - return m_value_cache_shape; + return m_value_cache_shape[id]; } size_t get_num_kv_blocks() const { @@ -153,7 +168,11 @@ class DeviceConfig { } size_t get_block_size_in_bytes() const { - return m_num_decoder_layers * 2 * m_num_kv_heads * m_block_size * m_head_size * get_cache_precision().size(); + size_t block_size = 0; + for (size_t layer_id = 0; layer_id < m_num_decoder_layers; layer_id++) { + block_size += 2 * m_num_kv_heads[layer_id] * m_block_size * m_head_size * get_cache_precision().size(); + } + return block_size; } }; } diff --git a/src/cpp/src/utils/paged_attention_transformations.cpp b/src/cpp/src/utils/paged_attention_transformations.cpp index 4dedcf989a..f564be8f19 100644 --- a/src/cpp/src/utils/paged_attention_transformations.cpp +++ b/src/cpp/src/utils/paged_attention_transformations.cpp @@ -53,15 +53,21 @@ void set_kv_cache_type_and_shape(std::shared_ptr model, DeviceConfig& OPENVINO_ASSERT(key_cache_params.count(key_cache_param_name) != 0, "key_cache.0 tensor not found among model parameters"); ov::PartialShape k_shape = key_cache_params[key_cache_param_name]->get_partial_shape(); OPENVINO_ASSERT(k_shape.rank().get_length() == 3, "KV cache shape is expected to have rank 3, while shape is ", k_shape); - size_t num_kv_heads = k_shape[1].get_length(), head_size = k_shape[2].get_length(); - + size_t head_size = k_shape[2].get_length(); + std::vector num_kv_heads(num_layers); + for (size_t idx = 0; idx < num_layers; idx++) { + size_t num_heads = key_cache_params[std::string("key_cache.") + std::to_string(idx)]->get_partial_shape()[1].get_length(); + num_kv_heads[idx] = num_heads; + } device_config.set_model_params(num_kv_heads, head_size, num_layers); - for (auto it_k = key_cache_params.begin(), it_v = value_cache_params.begin(); it_k != key_cache_params.end();++it_k, ++it_v) { - it_k->second->set_element_type(device_config.get_cache_precision()); - it_v->second->set_element_type(device_config.get_cache_precision()); - it_k->second->set_partial_shape(device_config.get_key_cache_shape()); - it_v->second->set_partial_shape(device_config.get_value_cache_shape()); + for (size_t idx = 0; idx < num_layers; idx++) { + auto k = key_cache_params[std::string("key_cache.") + std::to_string(idx)]; + auto v = value_cache_params[std::string("value_cache.") + std::to_string(idx)]; + k->set_element_type(device_config.get_cache_precision()); + v->set_element_type(device_config.get_cache_precision()); + k->set_partial_shape(device_config.get_key_cache_shape(idx)); + v->set_partial_shape(device_config.get_value_cache_shape(idx)); } model->validate_nodes_and_infer_types(); diff --git a/tests/cpp/cache_manager.cpp b/tests/cpp/cache_manager.cpp index 7f07980389..5dc848aba5 100644 --- a/tests/cpp/cache_manager.cpp +++ b/tests/cpp/cache_manager.cpp @@ -54,7 +54,8 @@ TEST(TestCacheManager, test_cache_size_param) { const std::string device = "CPU"; ov::genai::DeviceConfig device_config(core, scheduler_config, "CPU"); size_t num_decoder_layers = 12; - device_config.set_model_params(12, 64, num_decoder_layers); + std::vector num_kv_heads(12, 12); + device_config.set_model_params(num_kv_heads, 64, num_decoder_layers); ov::InferRequest request = core.compile_model(get_dummy_model(num_decoder_layers)).create_infer_request(); auto cache_manager = std::make_shared(device_config, request, core); @@ -76,7 +77,8 @@ TEST(TestCacheManager, test_kv_blocks_param) { const std::string device = "CPU"; ov::genai::DeviceConfig device_config(core, scheduler_config, "CPU"); size_t num_decoder_layers = 12; - device_config.set_model_params(12, 64, num_decoder_layers); + std::vector num_kv_heads(12, 12); + device_config.set_model_params(num_kv_heads, 64, num_decoder_layers); ov::InferRequest request = core.compile_model(get_dummy_model(num_decoder_layers)).create_infer_request(); auto cache_manager = std::make_shared(device_config, request, core); @@ -97,9 +99,12 @@ TEST(TestCacheManager, test_dynamic_cache_increase) { ov::genai::DeviceConfig device_config(core, scheduler_config, "CPU"); size_t num_decoder_layers = 12; size_t head_size = 64; - size_t num_kv_heads = 12; + std::vector num_kv_heads(12, 12); device_config.set_model_params(num_kv_heads, head_size, num_decoder_layers); - size_t block_size_in_bytes = num_decoder_layers * 2 * num_kv_heads * device_config.get_block_size() * head_size * device_config.get_cache_precision().size(); + size_t block_size_in_bytes = 0; + for (size_t layer_id = 0; layer_id < num_decoder_layers; layer_id++) { + block_size_in_bytes += 2 * num_kv_heads[layer_id] * device_config.get_block_size() * head_size * device_config.get_cache_precision().size(); + } ov::InferRequest request = core.compile_model(get_dummy_model(num_decoder_layers)).create_infer_request(); diff --git a/tests/cpp/device_config.cpp b/tests/cpp/device_config.cpp index 0d7435818f..973648f637 100644 --- a/tests/cpp/device_config.cpp +++ b/tests/cpp/device_config.cpp @@ -18,7 +18,7 @@ TEST(TestDeviceConfig, kv_cache_precision_u8) { const std::string device = "CPU"; size_t num_decoder_layers = 12; size_t head_size = 64, head_size_u8 = head_size + 8; - size_t num_kv_heads = 12; + std::vector num_kv_heads(12, 12); ov::genai::DeviceConfig device_config_default(core, scheduler_config, "CPU"); device_config_default.set_model_params(num_kv_heads, head_size_u8, num_decoder_layers); diff --git a/tests/cpp/scheduler.cpp b/tests/cpp/scheduler.cpp index ea1720faa2..cc0b53a433 100644 --- a/tests/cpp/scheduler.cpp +++ b/tests/cpp/scheduler.cpp @@ -44,7 +44,7 @@ std::shared_ptr init_cache_manager(SchedulerConfig scheduler_confi size_t num_decoder_layers = 12; ov::InferRequest request = core.compile_model(get_model(num_decoder_layers)).create_infer_request(); size_t head_size = 64, head_size_u8 = head_size + 8; - size_t num_kv_heads = 12; + std::vector num_kv_heads(12, 12); ov::genai::DeviceConfig device_config(core, scheduler_config, "CPU"); device_config.set_model_params(num_kv_heads, head_size_u8, num_decoder_layers); return std::make_shared(device_config, request, core);