From 8c70b8fa19e3b43442f4bb54956c16ef9bb221cd Mon Sep 17 00:00:00 2001 From: Bowen Wang Date: Thu, 22 May 2025 09:41:21 +0000 Subject: [PATCH] [Bugfix] Use random hidden states in dummy sampler run This is to avoid special values like `inf` or `nan` breaking the sampler kernel. Signed-off-by: Bowen Wang --- vllm/v1/worker/gpu_model_runner.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e26f97d816ae..759d69293a32 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1721,6 +1721,10 @@ def _dummy_sampler_run( self, hidden_states: torch.Tensor, ) -> torch.Tensor: + # The dummy hidden states may contain special values, + # like `inf` or `nan`. + # To avoid breaking the sampler, we use a random tensor here instead. + hidden_states = torch.rand_like(hidden_states) logits = self.model.compute_logits(hidden_states, None) num_reqs = logits.size(0)