Merge branch 'main' into wentao-optimize-model-runner-v2-sampler

Signed-off-by: yewentao256 <zhyanwentao@126.com>

vllm/v1/worker/gpu/sample/bad_words.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import numpy as np
 import torch
 
 from vllm.sampling_params import SamplingParams
@@ -73,15 +72,10 @@ class BadWordsState:
         self,
         logits: torch.Tensor,
         expanded_idx_mapping: torch.Tensor,
-        idx_mapping_np: np.ndarray,
         input_ids: torch.Tensor,
         expanded_local_pos: torch.Tensor,
+        max_num_bad_words: int,
     ) -> None:
-        max_num_bad_words = int(self.num_bad_words.np[idx_mapping_np].max())
-        if max_num_bad_words == 0:
-            # No request uses bad words. Skip the kernel launch.
-            return
-
         apply_bad_words(
             logits,
             expanded_idx_mapping,

vllm/v1/worker/gpu/sample/logit_bias.py

@@ -122,13 +122,8 @@ class LogitBiasState:
         self,
         logits: torch.Tensor,
         expanded_idx_mapping: torch.Tensor,
-        idx_mapping_np: np.ndarray,
         pos: torch.Tensor,
     ) -> None:
-        if not np.any(self.use_logit_bias[idx_mapping_np]):
-            # No request uses logit bias. Skip the kernel launch.
-            return
-
         apply_logit_bias(
             logits,
             expanded_idx_mapping,

vllm/v1/worker/gpu/sample/penalties.py

@@ -82,14 +82,9 @@ class PenaltiesState:
         self,
         logits: torch.Tensor,
         expanded_idx_mapping: torch.Tensor,
-        idx_mapping_np: np.ndarray,
         input_ids: torch.Tensor,
         expanded_local_pos: torch.Tensor,
     ) -> None:
-        if not np.any(self.use_penalty[idx_mapping_np]):
-            # No request uses penalties. Skip the kernel launch.
-            return
-
         apply_penalties(
             logits,
             expanded_idx_mapping,

vllm/v1/worker/gpu/sample/sampler.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass
+
 import numpy as np
 import torch
 
@@ -22,6 +24,34 @@ from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS, SamplingStates
 from vllm.v1.worker.gpu.states import RequestState
 
 
+@dataclass(frozen=True, slots=True)
+class _LogitsProcessingPlan:
+    logit_bias: bool
+    penalties: bool
+    bad_words: bool
+    temperature: bool
+    min_p: bool
+    top_k: bool
+    top_p: bool
+    max_num_bad_words: int
+
+    @property
+    def requires_processing(self) -> bool:
+        return (
+            self.logit_bias
+            or self.penalties
+            or self.bad_words
+            or self.temperature
+            or self.min_p
+            or self.top_k
+            or self.top_p
+        )
+
+    @property
+    def top_k_top_p(self) -> bool:
+        return self.top_k or self.top_p
+
+
 class Sampler:
     def __init__(
         self,
@@ -62,6 +92,82 @@ class Sampler:
         self.bad_words_state.apply_staged_writes()
         self.logprob_token_ids_state.apply_staged_writes()
 
+    def _build_logits_processing_plan(
+        self, idx_mapping_np: np.ndarray
+    ) -> _LogitsProcessingPlan:
+        max_num_bad_words = int(
+            self.bad_words_state.num_bad_words.np[idx_mapping_np].max()
+        )
+        temperature = self.sampling_states.temperature.np[idx_mapping_np]
+        return _LogitsProcessingPlan(
+            logit_bias=np.any(self.logit_bias_state.use_logit_bias[idx_mapping_np]),
+            penalties=np.any(self.penalties_state.use_penalty[idx_mapping_np]),
+            bad_words=max_num_bad_words > 0,
+            temperature=not np.all((temperature == 0.0) | (temperature == 1.0)),
+            min_p=np.any(self.sampling_states.min_p.np[idx_mapping_np] != 0.0),
+            top_k=np.any(
+                self.sampling_states.top_k.np[idx_mapping_np]
+                != self.sampling_states.vocab_size
+            ),
+            top_p=np.any(self.sampling_states.top_p.np[idx_mapping_np] != 1.0),
+            max_num_bad_words=max_num_bad_words,
+        )
+
+    def _apply_logits_processing(
+        self,
+        logits: torch.Tensor,
+        expanded_idx_mapping: torch.Tensor,
+        plan: _LogitsProcessingPlan,
+        pos: torch.Tensor,
+        input_ids: torch.Tensor,
+        expanded_local_pos: torch.Tensor,
+    ) -> torch.Tensor:
+        # Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.
+        if plan.logit_bias:
+            self.logit_bias_state.apply_logit_bias(
+                logits,
+                expanded_idx_mapping,
+                pos,
+            )
+
+        # Apply penalties in place.
+        if plan.penalties:
+            self.penalties_state.apply_penalties(
+                logits,
+                expanded_idx_mapping,
+                input_ids,
+                expanded_local_pos,
+            )
+
+        # Apply bad words masking in place.
+        if plan.bad_words:
+            self.bad_words_state.apply_bad_words(
+                logits,
+                expanded_idx_mapping,
+                input_ids,
+                expanded_local_pos,
+                plan.max_num_bad_words,
+            )
+
+        # Apply temperature in place.
+        if plan.temperature:
+            self.sampling_states.apply_temperature(logits, expanded_idx_mapping)
+
+        # Apply min_p in place.
+        if plan.min_p:
+            self.sampling_states.apply_min_p(logits, expanded_idx_mapping)
+
+        # Apply top_k and/or top_p. This might or might not return a new tensor.
+        if plan.top_k_top_p:
+            logits = self.sampling_states.apply_top_k_top_p(
+                logits,
+                expanded_idx_mapping,
+                plan.top_k,
+                plan.top_p,
+            )
+
+        return logits
+
     def __call__(
         self,
         logits: torch.Tensor,
@@ -129,44 +235,19 @@ class Sampler:
         input_ids: torch.Tensor,
         expanded_local_pos: torch.Tensor,
     ) -> torch.Tensor:
-        # Copy logits to a new FP32 tensor.
-        logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
-
-        # Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.
-        self.logit_bias_state.apply_logit_bias(
-            logits, expanded_idx_mapping, idx_mapping_np, pos
-        )
-
-        # Apply penalties in place.
-        self.penalties_state.apply_penalties(
-            logits,
-            expanded_idx_mapping,
-            idx_mapping_np,
-            input_ids,
-            expanded_local_pos,
-        )
-
-        # Apply bad words masking in place.
-        self.bad_words_state.apply_bad_words(
-            logits,
-            expanded_idx_mapping,
-            idx_mapping_np,
-            input_ids,
-            expanded_local_pos,
-        )
-
-        # Apply temperature in place.
-        self.sampling_states.apply_temperature(
-            logits, expanded_idx_mapping, idx_mapping_np
-        )
-
-        # Apply min_p in place.
-        self.sampling_states.apply_min_p(logits, expanded_idx_mapping, idx_mapping_np)
-
-        # Apply top_k and/or top_p. This might or might not return a new tensor.
-        return self.sampling_states.apply_top_k_top_p(
-            logits, expanded_idx_mapping, idx_mapping_np
-        )
+        processing_plan = self._build_logits_processing_plan(idx_mapping_np)
+        if processing_plan.requires_processing:
+            # Copy logits to a new FP32 tensor.
+            logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
+            logits = self._apply_logits_processing(
+                logits,
+                expanded_idx_mapping,
+                processing_plan,
+                pos,
+                input_ids,
+                expanded_local_pos,
+            )
+        return logits
 
     def sample(
         self,

vllm/v1/worker/gpu/sample/states.py

@@ -65,34 +65,23 @@ class SamplingStates:
         self,
         logits: torch.Tensor,
         expanded_idx_mapping: torch.Tensor,
-        idx_mapping_np: np.ndarray,
     ) -> None:
-        temp_np = self.temperature.np[idx_mapping_np]
-        if np.all((temp_np == 0.0) | (temp_np == 1.0)):
-            # No request requires temperature. Skip the kernel launch.
-            return
-
         apply_temperature(logits, expanded_idx_mapping, self.temperature.gpu)
 
     def apply_min_p(
         self,
         logits: torch.Tensor,
         expanded_idx_mapping: torch.Tensor,
-        idx_mapping_np: np.ndarray,
     ) -> None:
-        if np.all(self.min_p.np[idx_mapping_np] == 0.0):
-            # No request uses min_p. Skip the kernel launch.
-            return
         apply_min_p(logits, expanded_idx_mapping, self.min_p.gpu)
 
     def apply_top_k_top_p(
         self,
         logits: torch.Tensor,
         expanded_idx_mapping: torch.Tensor,
-        idx_mapping_np: np.ndarray,
+        do_top_k: bool,
+        do_top_p: bool,
     ) -> torch.Tensor:
-        do_top_k = np.any(self.top_k.np[idx_mapping_np] != self.vocab_size)
-        do_top_p = np.any(self.top_p.np[idx_mapping_np] != 1.0)
         if not (do_top_k or do_top_p):
             return logits