xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-09 18:57:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

updated

Browse Source 
Signed-off-by: Robert Shaw <robshaw@redhat.com>
This commit is contained in:
Robert Shaw 2025-03-28 02:17:42 +00:00
parent c5d963835b
commit 24f68342b4
1 changed files with 33 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

49
vllm/v1/worker/tpu_model_runner.py
View File

@ -604,25 +604,30 @@ class TPUModelRunner:
# avoid recompilations. # avoid recompilations.
tpu_sampling_metadata = TPUSupportedSamplingMetadata.\ tpu_sampling_metadata = TPUSupportedSamplingMetadata.\
from_input_batch(self.input_batch, logits_indices) from_input_batch(self.input_batch, logits_indices)
# Run the decoder
with set_forward_context(attn_metadata, self.vllm_config): # Temporary debug pathway.
hidden_states = self.model(
input_ids=input_ids,
positions=self.position_ids,
kv_caches=self.kv_caches,
inputs_embeds=inputs_embeds,
)
if envs.VLLM_TPU_DISABLE_SAMPLER_DEBUG: if envs.VLLM_TPU_DISABLE_SAMPLER_DEBUG:
selected_token_ids = self.model.compute_logits(hidden_states, with set_forward_context(attn_metadata, self.vllm_config):
logits_indices, None) hidden_states = self.model(
input_ids=input_ids,
positions=self.position_ids,
kv_caches=self.kv_caches,
inputs_embeds=inputs_embeds,
)
selected_token_ids = self.model.compute_logits_no_sampler(
hidden_states, logits_indices, None)
selected_token_ids = selected_token_ids.cpu()[:num_reqs] selected_token_ids = selected_token_ids.cpu()[:num_reqs]
else: else:
selected_token_ids = self.model.sample_from_hidden( with set_forward_context(attn_metadata, self.vllm_config):
hidden_states, tpu_sampling_metadata) hidden_states = self.model(
input_ids=input_ids,
# Remove padding on cpu and keep dynamic op outside of xla graph. positions=self.position_ids,
selected_token_ids = selected_token_ids.cpu()[:num_reqs] kv_caches=self.kv_caches,
inputs_embeds=inputs_embeds,
)
selected_token_ids = self.model.sample_from_hidden(
hidden_states, tpu_sampling_metadata)
selected_token_ids = selected_token_ids.cpu()[:num_reqs]
# Update the cache state concurrently. Code above will not block until # Update the cache state concurrently. Code above will not block until
# we use `selected_token_ids`. Add mark_step if post-processing changes # we use `selected_token_ids`. Add mark_step if post-processing changes
@ -936,6 +941,18 @@ class ModelWrapperV1(nn.Module):
logits = self.model.compute_logits(hidden_states, None) logits = self.model.compute_logits(hidden_states, None)
return logits return logits
@torch.compile(backend="openxla", fullgraph=True, dynamic=False)
def compute_logits_no_sampler(
self,
hidden_states: torch.Tensor,
logits_indices: torch.Tensor,
sampling_metadata,
) -> Optional[torch.Tensor]:
hidden_states = hidden_states[logits_indices]
logits = self.model.compute_logits(hidden_states, sampling_metadata)
selected_token_ids = torch.argmax(logits, dim=-1, keepdim=True)
return selected_token_ids
def get_multimodal_embeddings(self, *args, **kwargs): def get_multimodal_embeddings(self, *args, **kwargs):
return self.model.get_multimodal_embeddings(*args, **kwargs) return self.model.get_multimodal_embeddings(*args, **kwargs)