Merge pull request vllm-project#15 from luo-cheng2021/luocheng/openvi…

…no-model-executor-opt [CPU] Avoid copy result and force allocation
slyalin · Mar 25, 2024 · f3a397f · f3a397f
2 parents 05b9161 + 0f83539
commit f3a397f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/vllm/executor/openvino_executor.py b/vllm/executor/openvino_executor.py
@@ -119,6 +119,9 @@ def allocate_cpu_cache(self) -> List[OpenVINOKVCache]:
         for _ in range(self.num_layers):
             key_blocks = ov.Tensor(self.cache_dtype, key_block_shape)
             value_blocks = ov.Tensor(self.cache_dtype, value_block_shape)
+            # force allocation
+            key_blocks.data[:] = 0
+            value_blocks.data[:] = 0
             cpu_cache.append((key_blocks, value_blocks))
         return cpu_cache
 

diff --git a/vllm/model_executor/openvino_model_loader.py b/vllm/model_executor/openvino_model_loader.py
@@ -51,8 +51,9 @@ def ov_wrapper(self, *args, **kwargs) -> torch.Tensor:
     else:
         inputs.append(np.array(0, dtype=np.int32))   # for optimum-based models this parameter can be used even on the first iteration
 
-    outputs = self._ov_request.infer(inputs, share_inputs=True, share_outputs=False)
-    return torch.from_numpy(outputs[0])
+    self._ov_request.start_async(inputs, share_inputs=True)
+    self._ov_request.wait()
+    return torch.from_numpy(self._ov_request.get_tensor("logits").data)
 
 
 def patch_stateful_model(