Skip to content

GitLab

Unverified Commit 07064cb1 authored by Lu Fang's avatar Lu Fang Committed by GitHub
Browse files

[Bugfix] Check chain_speculative_sampling before calling it (#11673) 


Signed-off-by: default avatarLu Fang <lufang@fb.com>
parent 2f1e8e8f
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
vllm/model_executor/layers/rejection_sampler.py
View file @ 07064cb1
...@@ -118,7 +118,7 @@ class RejectionSampler(SpecDecodeStochasticBaseSampler): ...@@ -118,7 +118,7 @@ class RejectionSampler(SpecDecodeStochasticBaseSampler):
# If use Flashinfer chain_speculative_sampling kernel # If use Flashinfer chain_speculative_sampling kernel
# for rejection sampling # for rejection sampling
if self.use_flashinfer: if self.use_flashinfer and chain_speculative_sampling is not None:
batch_size, k, _ = draft_probs.shape batch_size, k, _ = draft_probs.shape
uniform_samples = self._create_uniform_samples( uniform_samples = self._create_uniform_samples(
seeded_seqs, batch_size, k, draft_probs.device) seeded_seqs, batch_size, k, draft_probs.device)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment