Skip to content

Commit

Permalink
[GPU] Fix issue to calculate present layout's padding for KVCache (#2…
Browse files Browse the repository at this point in the history
…5682)

### Details:
 - Fix issue to calculate present layout's padding for KVCache

### Tickets:
 - 146876
  • Loading branch information
andrew-k-park authored Jul 23, 2024
1 parent d0c07cd commit bb5a9d4
Showing 1 changed file with 7 additions and 2 deletions.
9 changes: 7 additions & 2 deletions src/plugins/intel_gpu/src/graph/primitive_inst.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1189,6 +1189,7 @@ void primitive_inst::do_runtime_in_place_kv_cache() {
}
const auto& desc = _node->as<kv_cache>().get_primitive();
auto& past_layout = _impl_params->input_layouts[0];
auto& new_layout = _impl_params->input_layouts[1];
auto& present_layout = _impl_params->output_layouts[0];
const auto& sequence_axis = desc->concat_axis;
const auto& gather_axis = desc->gather_axis;
Expand All @@ -1209,8 +1210,12 @@ void primitive_inst::do_runtime_in_place_kv_cache() {
auto max_pad = kv_cache_inst::get_max_pad(past_layout, _deps[0].first->_max_output_layout_count[0], sequence_axis_legacy, "past_layout");

if (max_pad > 0) {
kv_cache_inst::update_pad(present_layout, max_pad - 1, sequence_axis_legacy);
GPU_DEBUG_TRACE_DETAIL << "[do runtime_in_place_kv_cache] " << id() << " Updated present_layout's pad : " << present_layout.to_string() << std::endl;
const auto new_seq_len = static_cast<int64_t>(new_layout.get_shape()[sequence_axis]);
if (max_pad - new_seq_len >= 0) {
kv_cache_inst::update_pad(present_layout, max_pad - new_seq_len, sequence_axis_legacy);
GPU_DEBUG_TRACE_DETAIL << "[do runtime_in_place_kv_cache] " << id() << " Updated present_layout's pad : "
<< present_layout.to_string() << std::endl;
}
auto& variable = get_network().get_variable(desc->variable_info.variable_id);
variable.set_layout(present_layout);
GPU_DEBUG_TRACE_DETAIL << "[do_runtime_in_place_kv_cache] " << id() << "Updated variable with present_layout"
Expand Down

0 comments on commit bb5a9d4

Please sign in to comment.