clean up

python/llm/example/NPU/HF-Transformers-AutoModels/LLM/generate.py

@@ -24,14 +24,14 @@
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Predict Tokens using `generate()` API for npu model')
-    parser.add_argument('--repo-id-or-model-path', type=str, default=r'C:\Users\intel\models\Qwen2-7B-Instruct',
+    parser.add_argument('--repo-id-or-model-path', type=str, default="D:\llm-models\Llama-2-7b-chat-hf",
                         help='The huggingface repo id for the Llama2 model to be downloaded'
                              ', or the path to the huggingface checkpoint folder')
     parser.add_argument('--prompt', type=str, default="Once upon a time, there existed a little girl who liked to have adventures. She wanted to go to places and meet new people, and have fun",
                         help='Prompt to infer')
     parser.add_argument('--n-predict', type=int, default=32,
                         help='Max tokens to predict')
-    parser.add_argument('--load_in_low_bit', type=str, default="sym_int4",
+    parser.add_argument('--load_in_low_bit', type=str, default="sym_int8",
                         help='Load in low bit to use')
 
     args = parser.parse_args()
@@ -45,7 +45,7 @@
     print(model)
 
     with torch.inference_mode():
-        prompt = "Once upon a time, there existed a little girl who liked to have adventures. She wanted to go to places and meet new people, and "
+        prompt = "Once upon a time, there existed a little girl who liked to have adventures. She wanted to go to places and meet new people, and have fun"
         input_ids = tokenizer.encode(prompt, return_tensors="pt")
         print("finish to load")
         print('input length:', len(input_ids[0]))

python/llm/setup.py

@@ -233,23 +233,23 @@ def setup_package():
     change_permission = True if platform_name == "Linux" else False
 
     # Delete legacy libs
-    # if os.path.exists(libs_dir):
-    #     print(f"Deleting existing libs_dir {libs_dir} ....")
-    #     shutil.rmtree(libs_dir)
-    # os.makedirs(libs_dir, exist_ok=True)
-    # open(os.path.join(libs_dir, "__init__.py"), 'w').close()
-
-    # # copy built files for github workflow
-    # for built_file in glob.glob(os.path.join(github_artifact_dir, '*')):
-    #     print(f'Copy workflow built file: {built_file}')
-    #     if change_permission:
-    #         os.chmod(built_file, 0o775)
-    #     shutil.copy(built_file, libs_dir)
-
-    # lib_urls = obtain_lib_urls()
-
-    # for url in lib_urls[platform_name]:
-    #     download_libs(url, change_permission=change_permission)
+    if os.path.exists(libs_dir):
+        print(f"Deleting existing libs_dir {libs_dir} ....")
+        shutil.rmtree(libs_dir)
+    os.makedirs(libs_dir, exist_ok=True)
+    open(os.path.join(libs_dir, "__init__.py"), 'w').close()
+
+    # copy built files for github workflow
+    for built_file in glob.glob(os.path.join(github_artifact_dir, '*')):
+        print(f'Copy workflow built file: {built_file}')
+        if change_permission:
+            os.chmod(built_file, 0o775)
+        shutil.copy(built_file, libs_dir)
+
+    lib_urls = obtain_lib_urls()
+
+    for url in lib_urls[platform_name]:
+        download_libs(url, change_permission=change_permission)
 
     # Check if all package files are ready
     for file in package_data[platform_name]:

python/llm/src/ipex_llm/transformers/kv.py

@@ -101,9 +101,8 @@ def update(
                 batch_size, num_heads, head_dim,
                 0, max_len,
                 key_states.dtype, key_states.device,
-                tranpose_value=transpose_value
             )
-            k_cache, v_cache = append_kv_cache(k_cache, v_cache, key_states, value_states, transpose_value=transpose_value)
+            k_cache, v_cache = append_kv_cache(k_cache, v_cache, key_states, value_states)
 
             self.key_cache[layer_idx] = k_cache
             self.value_cache[layer_idx] = v_cache
@@ -123,7 +122,7 @@ def update(
                 new_v_cache[...] = v_cache[...]
                 k_cache = new_k_cache
                 v_cache = new_v_cache
-            k_cache, v_cache = append_kv_cache(k_cache, v_cache, key_states, value_states, transpose_value=transpose_value)
+            k_cache, v_cache = append_kv_cache(k_cache, v_cache, key_states, value_states)
             self.key_cache[layer_idx] = k_cache
             self.value_cache[layer_idx] = v_cache
 

python/llm/src/ipex_llm/transformers/models/utils.py

@@ -37,41 +37,23 @@ def decoding_fast_path_qtype_check(proj):
     return qtype in [SYM_INT4, FP8E5, FP4]
 
 
-def init_kv_cache(batch_size, num_heads, head_dim, current_length, max_length, dtype, device, tranpose_value=False):
-    if not tranpose_value:
-        key_cache_storage = torch.zeros(batch_size, num_heads,
-                                        max_length, head_dim,
-                                        dtype=dtype, device=device)
-        value_cache_storage = torch.zeros(batch_size, num_heads,
-                                        max_length, head_dim,
-                                        dtype=dtype, device=device)
-
-        key_cache = key_cache_storage.as_strided((batch_size, num_heads,
+def init_kv_cache(batch_size, num_heads, head_dim, current_length, max_length, dtype, device):
+    key_cache_storage = torch.zeros(batch_size, num_heads,
+                                    max_length, head_dim,
+                                    dtype=dtype, device=device)
+    value_cache_storage = torch.zeros(batch_size, num_heads,
+                                    max_length, head_dim,
+                                    dtype=dtype, device=device)
+
+    key_cache = key_cache_storage.as_strided((batch_size, num_heads,
+                                            current_length, head_dim),
+                                            key_cache_storage.stride(),
+                                            storage_offset=0)
+    value_cache = value_cache_storage.as_strided((batch_size, num_heads,
                                                 current_length, head_dim),
-                                                key_cache_storage.stride(),
+                                                value_cache_storage.stride(),
                                                 storage_offset=0)
-        value_cache = value_cache_storage.as_strided((batch_size, num_heads,
-                                                    current_length, head_dim),
-                                                    value_cache_storage.stride(),
-                                                    storage_offset=0)
-        return key_cache, value_cache
-    else:
-        key_cache_storage = torch.zeros(batch_size, num_heads,
-                                        max_length, head_dim,
-                                        dtype=dtype, device=device)
-        value_cache_storage = torch.zeros(batch_size, num_heads,
-                                          head_dim, max_length,
-                                          dtype=dtype, device=device)
-
-        key_cache = key_cache_storage.as_strided((batch_size, num_heads,
-                                                current_length, head_dim),
-                                                key_cache_storage.stride(),
-                                                storage_offset=0)
-        value_cache = value_cache_storage.as_strided((batch_size, num_heads,
-                                                    head_dim, current_length),
-                                                    value_cache_storage.stride(),
-                                                    storage_offset=0)
-        return key_cache, value_cache.transpose(-1, -2)
+    return key_cache, value_cache
 
 
 def extend_kv_cache(batch_size, num_heads, head_dim, current_length, max_length, dtype, device):
@@ -81,35 +63,16 @@ def extend_kv_cache(batch_size, num_heads, head_dim, current_length, max_length,
     return init_kv_cache(batch_size, num_heads, head_dim, current_length, max_length, dtype, device)
 
 
-def append_kv_cache(cache_k, cache_v, key_states, value_states, transpose_value=False):
-    if not transpose_value:
-        new_size = (cache_k.size(0),
-                    cache_k.size(1),
-                    cache_k.size(2) + key_states.size(2),
-                    cache_k.size(3))
-        new_cache_k = cache_k.as_strided(new_size, cache_k.stride(), storage_offset=0)
-        new_cache_k[:, :, cache_k.size(2):cache_k.size(2) + key_states.size(2), :] = key_states
-        new_cache_v = cache_v.as_strided(new_size, cache_v.stride(), storage_offset=0)
-        new_cache_v[:, :, cache_v.size(2):cache_v.size(2) + key_states.size(2), :] = value_states
-        return new_cache_k, new_cache_v
-    else:
-        new_size_key = (cache_k.size(0),
-                        cache_k.size(1),
-                        cache_k.size(2) + key_states.size(2),
-                        cache_k.size(3))
-        new_cache_k = cache_k.as_strided(new_size_key, cache_k.stride(), storage_offset=0)
-        new_cache_k[:, :, cache_k.size(2):cache_k.size(2) + key_states.size(2), :] = key_states
-
-        new_size_value = (cache_v.size(0),
-                          cache_v.size(1),
-                          cache_v.size(3),
-                          cache_v.size(2) + value_states.size(3),
-                          )
-        raw_cache_v = cache_v.transpose(-1, -2)
-        # assert raw_cache_v.is_contiguous(), f"raw_cache_v size is {raw_cache_v.shape}, stride is {raw_cache_v.stride()}"
-        new_cache_v = raw_cache_v.as_strided(new_size_value, raw_cache_v.stride(), storage_offset=0)
-        new_cache_v[:, :, :, raw_cache_v.size(3):raw_cache_v.size(3) + value_states.size(3)] = value_states
-        return new_cache_k, new_cache_v.transpose(-1, -2)
+def append_kv_cache(cache_k, cache_v, key_states, value_states):
+    new_size = (cache_k.size(0),
+                cache_k.size(1),
+                cache_k.size(2) + key_states.size(2),
+                cache_k.size(3))
+    new_cache_k = cache_k.as_strided(new_size, cache_k.stride(), storage_offset=0)
+    new_cache_k[:, :, cache_k.size(2):cache_k.size(2) + key_states.size(2), :] = key_states
+    new_cache_v = cache_v.as_strided(new_size, cache_v.stride(), storage_offset=0)
+    new_cache_v[:, :, cache_v.size(2):cache_v.size(2) + key_states.size(2), :] = value_states
+    return new_cache_k, new_cache_v
 
 
 def use_quantize_kv_cache(linear: torch.nn.Module, x: torch.Tensor) -> bool: