[GPU] Fixed unexpected kv cache optimization turned off (#24805)

### Details:
- kv cache opt was disabled due to the wrongly detected beam table size.
Fixed this issue.
### Tickets:
 - *ticket-id*
This commit is contained in:
Taylor Yeonbok Lee 2024-06-03 03:34:13 -07:00 committed by GitHub
parent 142114ccbc
commit cea8e1f5e4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions

View File

@ -1034,8 +1034,8 @@ void primitive_inst::do_runtime_in_place_kv_cache() {
const auto& sequence_axis = desc->concat_axis;
const auto& gather_axis = desc->gather_axis;
const auto& prev_batch_size = past_layout.get_shape()[gather_axis];
const auto& beam_size = present_layout.get_shape()[gather_axis];
const auto& prev_batch_size = static_cast<size_t>(past_layout.get_shape()[gather_axis]);
const auto& beam_size = static_cast<size_t>(present_layout.get_shape()[gather_axis]);
if (prev_batch_size != beam_size) {
// If the previous batch size is not same as beam size, need explicit concat
_impl_params->_can_be_optimized = false;