update

Author: hkvision

python/llm/src/ipex_llm/transformers/npu_models/convert.py

@@ -404,8 +404,8 @@ def optimize_llm_single_process(
         invalidInputError(False,
                           "False to InitLLMPipeline.")
     # patch generate function
-    # import types
-    # model.generate = types.MethodType(generate, model)
+    import types
+    model.simple_generate = types.MethodType(generate, model)
     from transformers.modeling_utils import PreTrainedModel
     general_convert(model, PreTrainedModel, prepare_input_ids, "prepare_inputs_for_generation")
     general_convert(model, PreTrainedModel, causal_lm_forward)
@@ -439,17 +439,16 @@ def causal_lm_forward(
     return_dict: Optional[bool] = None,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     start = time.perf_counter()
-    from .npu_llm_cpp import run_decode, run_prefill, get_logits
+    from .npu_llm_cpp import run_prefill_with_logits, run_decode_with_logits
     if isinstance(input_ids[0], torch.Tensor):
         input_list = input_ids[0].flatten().tolist()
     else:
         input_list = input_ids[0]
     input_length = len(input_list)
     if input_length > 1:
-        run_prefill(self.model_ptr, input_list, self.vocab_size)
+        logits = run_prefill_with_logits(self.model_ptr, input_list, self.logits_buffer, self.vocab_size)
     else:
-        run_decode(self.model_ptr, input_list[0], self.vocab_size)
-    logits = get_logits(self.model_ptr, self.logits_buffer)
+        logits = run_decode_with_logits(self.model_ptr, input_list[0], self.logits_buffer, self.vocab_size)
     end = time.perf_counter()
     overall = (end - start) * 1000
     print("Overall time: ", overall)

python/llm/src/ipex_llm/transformers/npu_models/npu_llm_cpp.py

@@ -60,8 +60,11 @@ def get_shared_lib_info(lib_base_name: str):
 _lib.reset.argtypes = [ctypes.c_void_p]
 _lib.reset.restype = None
 
-_lib.get_logits.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_float)]
-_lib.reset.restype = None
+_lib.run_prefill_with_logits.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.POINTER(ctypes.c_float), ctypes.c_int]
+_lib.run_prefill_with_logits.restype = None
+
+_lib.run_decode_with_logits.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.POINTER(ctypes.c_float), ctypes.c_int]
+_lib.run_decode_with_logits.restype = None
 
 
 def load_model_from_file(model_dir: str):
@@ -82,12 +85,21 @@ def run_decode(model_ptr, input_id, vocab_size):
     return new_token
 
 
-def reset(model_ptr):
-    _lib.reset(model_ptr)
+def run_prefill_with_logits(model_ptr, input_ids, logits, vocab_size):
+    input_ptr = (ctypes.c_int32 * len(input_ids))(*input_ids)
+    input_len = len(input_ids)
+    logits_ptr = logits.data.data_ptr()
+    logits_ptr = ctypes.cast(logits_ptr, ctypes.POINTER(ctypes.c_float))
+    _lib.run_prefill_with_logits(model_ptr, input_ptr, input_len, logits_ptr, vocab_size)
+    return logits
 
 
-def get_logits(model_ptr, logits):
-    src = logits.data.data_ptr()
-    src = ctypes.cast(src, ctypes.POINTER(ctypes.c_float))
-    _lib.get_logits(model_ptr, src)
+def run_decode_with_logits(model_ptr, input_id, logits, vocab_size):
+    logits_ptr = logits.data.data_ptr()
+    logits_ptr = ctypes.cast(logits_ptr, ctypes.POINTER(ctypes.c_float))
+    _lib.run_decode_with_logits(model_ptr, input_id, logits_ptr, vocab_size)
     return logits
+
+
+def reset(model_ptr):
+    _lib.reset(model_ptr)