xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-10 01:05:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Minor][Models] Fix Return Types of Llama & Eagle (#17220)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
Woosuk Kwon 2025-04-25 21:54:47 -07:00 committed by GitHub
parent 7bd0c7745c
commit b278911229
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
vllm/model_executor/models/llama.py
View File

@ -345,7 +345,8 @@ class LlamaModel(nn.Module):
positions: torch.Tensor, positions: torch.Tensor,
intermediate_tensors: Optional[IntermediateTensors], intermediate_tensors: Optional[IntermediateTensors],
inputs_embeds: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, IntermediateTensors]: ) -> Union[torch.Tensor, IntermediateTensors, tuple[torch.Tensor,
list[torch.Tensor]]]:
if get_pp_group().is_first_rank: if get_pp_group().is_first_rank:
if inputs_embeds is not None: if inputs_embeds is not None:
hidden_states = inputs_embeds hidden_states = inputs_embeds

4
vllm/model_executor/models/llama_eagle.py
View File

@ -70,7 +70,7 @@ class LlamaModel(nn.Module):
input_ids: torch.Tensor, input_ids: torch.Tensor,
positions: torch.Tensor, positions: torch.Tensor,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
) -> torch.Tensor: ) -> tuple[torch.Tensor, torch.Tensor]:
input_embeds = self.embed_tokens(input_ids) input_embeds = self.embed_tokens(input_ids)
hidden_states = self.fc( hidden_states = self.fc(
torch.cat((input_embeds, hidden_states), dim=-1)) torch.cat((input_embeds, hidden_states), dim=-1))
@ -133,7 +133,7 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):
input_ids: torch.Tensor, input_ids: torch.Tensor,
positions: torch.Tensor, positions: torch.Tensor,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
) -> torch.Tensor: ) -> tuple[torch.Tensor, torch.Tensor]:
return self.model(input_ids, positions, hidden_states) return self.model(input_ids, positions, hidden_states)
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):

4
vllm/model_executor/models/llama_eagle3.py
View File

@ -117,7 +117,7 @@ class LlamaModel(nn.Module):
input_ids: torch.Tensor, input_ids: torch.Tensor,
positions: torch.Tensor, positions: torch.Tensor,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
) -> torch.Tensor: ) -> tuple[torch.Tensor, torch.Tensor]:
input_embeds = self.embed_tokens(input_ids) input_embeds = self.embed_tokens(input_ids)
if (hidden_states.shape[-1] != input_embeds.shape[-1]): if (hidden_states.shape[-1] != input_embeds.shape[-1]):
hidden_states = self.fc(hidden_states) hidden_states = self.fc(hidden_states)
@ -194,7 +194,7 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
input_ids: torch.Tensor, input_ids: torch.Tensor,
positions: torch.Tensor, positions: torch.Tensor,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
) -> torch.Tensor: ) -> tuple[torch.Tensor, torch.Tensor]:
return self.model(input_ids, positions, hidden_states) return self.model(input_ids, positions, hidden_states)
def compute_logits( def compute_logits(