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, 22:15 UTC
Fix failing `MyGemma2Embedding` test (#13820)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Feb 25, 2025 at 20:33 UTC
34e3494e70a155bdfcec240c49a9b9caf0e28bba
1 file changed
+1
-6
tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_gemma_embedding.py
+1
-6
@@ -1,11 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-from typing import Iterable, List, Optional, Tuple, Union
3
+from typing import Iterable, Optional, Tuple, Union
4
5
import torch
6
import torch.nn as nn
7
8
-from vllm.attention import AttentionMetadata
8
from vllm.config import VllmConfig
9
from vllm.model_executor.layers.pooler import Pooler, PoolingType
10
from vllm.model_executor.models.gemma2 import Gemma2Model
@@ -37,16 +36,12 @@ class MyGemma2Embedding(nn.Module):
36
self,
37
input_ids: torch.Tensor,
38
positions: torch.Tensor,
40
- kv_caches: List[torch.Tensor],
41
- attn_metadata: AttentionMetadata,
39
intermediate_tensors: Optional[IntermediateTensors] = None,
40
inputs_embeds: Optional[torch.Tensor] = None,
41
) -> Union[torch.Tensor, IntermediateTensors]:
42
hidden_states = self.model(
43
input_ids,
44
positions,
48
- kv_caches,
49
- attn_metadata,
45
intermediate_tensors=intermediate_tensors,
46
inputs_embeds=inputs_embeds,
47
)