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

sample

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
Woosuk Kwon 2025-09-18 17:39:43 -07:00
parent 8deedfa42b
commit 52ca2f517a
1 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
vllm/v1/worker/gpu/model_runner.py
View File

@ -316,9 +316,12 @@ class GPUModelRunner:
def sample( def sample(
self, self,
logits: torch.Tensor, hidden_states: torch.Tensor,
input_batch: InputBatch, input_batch: InputBatch,
) -> SamplerOutput: ) -> SamplerOutput:
# TODO(woosuk): Support DP sampler + CUDA graphs.
sample_hidden_states = hidden_states[input_batch.logits_indices]
logits = self.model.compute_logits(sample_hidden_states, None)
pos = input_batch.positions[input_batch.logits_indices] pos = input_batch.positions[input_batch.logits_indices]
sampling_metadata = self.req_states.make_sampling_metadata( sampling_metadata = self.req_states.make_sampling_metadata(
input_batch.idx_mapping_np, pos) input_batch.idx_mapping_np, pos)
@ -369,11 +372,8 @@ class GPUModelRunner:
positions=input_batch.positions[:num_tokens], positions=input_batch.positions[:num_tokens],
) )
# Compute logits to sample next tokens. sampler_output = self.sample(hidden_states, input_batch)
sample_hidden_states = hidden_states[input_batch.logits_indices]
logits = self.model.compute_logits(sample_hidden_states, None)
sampler_output = self.sample(logits, input_batch)
sampled_token_ids_np, num_sampled_tokens = self.postprocess( sampled_token_ids_np, num_sampled_tokens = self.postprocess(
sampler_output, input_batch) sampler_output, input_batch)
logprobs = None logprobs = None