A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
Fix missing `kv_caches` and `attn_metadata` in `OpenVINOCausalLM` (#14271)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Mar 7, 2025 at 13:30 UTC
f7a6bd0fa1d9ede66c935c60017c494d1a852fc6
2 files changed
+11
-10
vllm/model_executor/model_loader/openvino.py
+9
-10
@@ -2,7 +2,7 @@
2
3
# ruff: noqa: SIM117
4
from pathlib import Path
5
-from typing import List, Optional, Tuple
5
+from typing import Optional
6
7
import openvino as ov
8
import torch
@@ -12,8 +12,8 @@ from optimum.intel import OVModelForCausalLM
12
from torch import nn
13
14
import vllm.envs as envs
15
-from vllm.attention.backends.openvino import OpenVINOAttentionMetadata
15
from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
16
+from vllm.forward_context import get_forward_context
17
from vllm.logger import init_logger
18
from vllm.model_executor.layers.logits_processor import (LogitsProcessor,
19
_prune_hidden_states)
@@ -24,7 +24,7 @@ from vllm.platforms import current_platform
24
logger = init_logger(__name__)
25
26
27
-def _flattenize_inputs(inputs):
27
+def _flatten_inputs(inputs):
28
"""
29
Helper function for making nested inputs flattens
30
"""
@@ -33,10 +33,9 @@ def _flattenize_inputs(inputs):
33
if input_data is None:
34
continue
35
if isinstance(input_data, (list, tuple)):
36
- flatten_inputs.extend(_flattenize_inputs(input_data))
36
+ flatten_inputs.extend(_flatten_inputs(input_data))
37
elif isinstance(input_data, dict):
38
- flatten_inputs.extend(_flattenize_inputs(list(
39
- input_data.values())))
38
+ flatten_inputs.extend(_flatten_inputs(list(input_data.values())))
39
else:
40
flatten_inputs.append(input_data)
41
return flatten_inputs
@@ -147,15 +146,15 @@ class OpenVINOCausalLM(nn.Module):
146
self,
147
input_ids: torch.Tensor,
148
positions: torch.Tensor,
150
- kv_caches: List[Tuple[ov.Tensor, ov.Tensor]],
151
- attn_metadata: OpenVINOAttentionMetadata,
149
+ kv_caches: list[tuple[ov.Tensor, ov.Tensor]],
150
) -> torch.Tensor:
153
- flatten_kv_cache = _flattenize_inputs(kv_caches)
151
+ flat_kv_caches = _flatten_inputs(kv_caches)
152
+ attn_metadata = get_forward_context().attn_metadata
153
154
inputs = [
155
input_ids,
156
positions,
158
- *flatten_kv_cache,
157
+ *flat_kv_caches,
158
attn_metadata.past_lens,
159
attn_metadata.subsequence_begins,
160
attn_metadata.block_indices,
vllm/worker/openvino_model_runner.py
+2
@@ -346,6 +346,8 @@ class OpenVINOModelRunner(ModelRunnerBase):
346
input_tokens,
347
"positions":
348
input_positions,
349
+ "kv_caches":
350
+ kv_caches,
351
**MultiModalKwargs.as_kwargs(multi_modal_kwargs or {},
352
device=self.device),
353
}