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

Fix some typo (#19475)

Browse Source 
Signed-off-by: ximing.wxm <ximing.wxm@antgroup.com>
Co-authored-by: ximing.wxm <ximing.wxm@antgroup.com>
This commit is contained in:
Ximingwang-09 2025-06-11 18:36:04 +08:00 committed by GitHub
parent 7484e1fce2
commit 3c8694eabe
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/model_executor/layers/rejection_sampler.py
View File

@ -283,14 +283,14 @@ class RejectionSampler(SpecDecodeStochasticBaseSampler):
batch_size, k, _ = draft_probs.shape batch_size, k, _ = draft_probs.shape
batch_indices = torch.arange(batch_size, batch_indices = torch.arange(batch_size,
device=target_probs.device)[:, None] device=target_probs.device)[:, None]
probs_indicies = torch.arange(k, device=target_probs.device) probs_indices = torch.arange(k, device=target_probs.device)
# shape [batch_size, k] # shape [batch_size, k]
selected_draft_probs = draft_probs[batch_indices, probs_indicies, selected_draft_probs = draft_probs[batch_indices, probs_indices,
draft_token_ids] draft_token_ids]
# shape [batch_size, k] # shape [batch_size, k]
selected_target_probs = target_probs[batch_indices, probs_indicies, selected_target_probs = target_probs[batch_indices, probs_indices,
draft_token_ids] draft_token_ids]
uniform_rand = self._create_uniform_samples(seeded_seqs, batch_size, uniform_rand = self._create_uniform_samples(seeded_seqs, batch_size,