[Bugfix] Check chain_speculative_sampling before calling it (vllm-pro…

…ject#11673) Signed-off-by: Lu Fang <lufang@fb.com>
neuralmagic · Jan 3, 2025 · 07064cb · 07064cb
1 parent 2f1e8e8
commit 07064cb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vllm/model_executor/layers/rejection_sampler.py b/vllm/model_executor/layers/rejection_sampler.py
@@ -118,7 +118,7 @@ def forward(
 
         # If use Flashinfer chain_speculative_sampling kernel
         # for rejection sampling
-        if self.use_flashinfer:
+        if self.use_flashinfer and chain_speculative_sampling is not None:
             batch_size, k, _ = draft_probs.shape
             uniform_samples = self._create_uniform_samples(
                 seeded_seqs, batch_size, k, draft_probs.device)