update gpu block size based on xattn

src/cpp/src/continuous_batching/cache_manager.hpp

@@ -45,8 +45,7 @@ class CacheManager {
                         "Continuous batching: execution device is expected to be single CPU / single GPU / multi GPUs");
         m_device = execution_devices[0];
         // set block_size depending on device
-        const size_t cpu_block_size = 32, gpu_block_size = 16;
-        m_block_size = all_gpu_device ? gpu_block_size : cpu_block_size;
+        const size_t cpu_block_size = 32, gpu_block_size = 16, gpu_block_size_xattn = 256;
 
         if (all_gpu_device) {
             m_context = m_request.get_compiled_model().get_context();
@@ -75,6 +74,13 @@ class CacheManager {
             }
         }
 
+        bool has_xattention = false;
+        if ((m_key_shapes[0][2].get_length() == m_value_shapes[0][2].get_length()) &&
+            (m_key_shapes[0][3].get_length() == m_value_shapes[0][3].get_length()) &&
+            (m_key_shapes[0][2].get_length() == gpu_block_size_xattn)) {
+            has_xattention = true;        
+        }
+        m_block_size = all_gpu_device ? ( has_xattention ? gpu_block_size_xattn : gpu_block_size ) : cpu_block_size;
         m_num_decoder_layers = m_value_precisions.size();
         OPENVINO_ASSERT(m_num_decoder_layers == m_key_precisions.size(), "Invalid case: a different number of K and V caches in a LLM model");
     }

tools/llm_bench/llm_bench_utils/ov_utils.py

@@ -190,10 +190,14 @@ def get_scheduler_config_genai(config_data, config_name="CB config"):
             sparse_attention_kwargs = user_config.pop('sparse_attention_config')
             if "mode" in sparse_attention_kwargs.keys():
                 sparse_attention_kwargs["mode"] = getattr(openvino_genai.SparseAttentionMode, sparse_attention_kwargs["mode"])
-
-            scheduler_config.use_sparse_attention = True
-            scheduler_config.sparse_attention_config = openvino_genai.SparseAttentionConfig(**sparse_attention_kwargs)
-            log.info("Sparse Attention mode ON")
+            if 'use_sparse_attention' in user_config.keys():
+                use_sparse_attention = user_config.pop('use_sparse_attention')
+            else:
+                use_sparse_attention = None
+            if use_sparse_attention != False:
+                scheduler_config.use_sparse_attention = True
+                scheduler_config.sparse_attention_config = openvino_genai.SparseAttentionConfig(**sparse_attention_kwargs)
+                log.info("Sparse Attention mode ON")
 
         for param, value in user_config.items():
             setattr(scheduler_config, param, value)

tools/llm_bench/task/text_generation.py

@@ -315,6 +315,7 @@ def run_text_generation_genai(input_text, num, model, tokenizer, args, iter_data
         if args['infer_count'] is not None:
             out_str += 'all max_output_token_size: {} * {}'.format(args['infer_count'], args['batch_size'])
         log.info(out_str)
+    from openvino_genai import GenerationConfig
     gen_config = model.get_generation_config() if hasattr(model, 'get_generation_config') else GenerationConfig()
     gen_config.max_new_tokens = max_gen_tokens
     # llama-3-8b-instruct's generation_config.json has 4096 max_length.