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, 13:30 UTC
[Minor][Models] Fix Return Types of Llama & Eagle (#17220)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Apr 25, 2025 at 21:54 UTC
b2789112294cb1f54c55c2011d2a5737c2b22cae
3 files changed
+6
-5
vllm/model_executor/models/llama.py
+2
-1
@@ -345,7 +345,8 @@ class LlamaModel(nn.Module):
345
positions: torch.Tensor,
346
intermediate_tensors: Optional[IntermediateTensors],
347
inputs_embeds: Optional[torch.Tensor] = None,
348
- ) -> Union[torch.Tensor, IntermediateTensors]:
348
+ ) -> Union[torch.Tensor, IntermediateTensors, tuple[torch.Tensor,
349
+ list[torch.Tensor]]]:
350
if get_pp_group().is_first_rank:
351
if inputs_embeds is not None:
352
hidden_states = inputs_embeds
vllm/model_executor/models/llama_eagle.py
+2
-2
@@ -70,7 +70,7 @@ class LlamaModel(nn.Module):
70
input_ids: torch.Tensor,
71
positions: torch.Tensor,
72
hidden_states: torch.Tensor,
73
- ) -> torch.Tensor:
73
+ ) -> tuple[torch.Tensor, torch.Tensor]:
74
input_embeds = self.embed_tokens(input_ids)
75
hidden_states = self.fc(
76
torch.cat((input_embeds, hidden_states), dim=-1))
@@ -133,7 +133,7 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):
133
input_ids: torch.Tensor,
134
positions: torch.Tensor,
135
hidden_states: torch.Tensor,
136
- ) -> torch.Tensor:
136
+ ) -> tuple[torch.Tensor, torch.Tensor]:
137
return self.model(input_ids, positions, hidden_states)
138
139
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
vllm/model_executor/models/llama_eagle3.py
+2
-2
@@ -117,7 +117,7 @@ class LlamaModel(nn.Module):
117
input_ids: torch.Tensor,
118
positions: torch.Tensor,
119
hidden_states: torch.Tensor,
120
- ) -> torch.Tensor:
120
+ ) -> tuple[torch.Tensor, torch.Tensor]:
121
input_embeds = self.embed_tokens(input_ids)
122
if (hidden_states.shape[-1] != input_embeds.shape[-1]):
123
hidden_states = self.fc(hidden_states)
@@ -194,7 +194,7 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
194
input_ids: torch.Tensor,
195
positions: torch.Tensor,
196
hidden_states: torch.Tensor,
197
- ) -> torch.Tensor:
197
+ ) -> tuple[torch.Tensor, torch.Tensor]:
198
return self.model(input_ids, positions, hidden_states)
199
200
def compute_logits(