A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 05:00 UTC
[Model Runner V2] Minor simplification for BadWordsState (#34669)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 16, 2026 at 21:27 UTC
9752da9d9c17f85bddee2f869130e72961f4d54d
1 file changed
+4
-12
vllm/v1/worker/gpu/sample/bad_words.py
+4
-12
@@ -39,18 +39,11 @@ class BadWordsState:
39
)
40
# number of bad words per request
41
self.num_bad_words = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
42
- # whether request uses bad words
43
- self.use_bad_words = np.zeros(self.max_num_reqs, dtype=bool)
42
45
- def add_request(
46
- self,
47
- req_idx: int,
48
- sampling_params: SamplingParams,
49
- ) -> None:
43
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
44
bad_words_token_ids = sampling_params.bad_words_token_ids
45
if not bad_words_token_ids:
46
self.num_bad_words.np[req_idx] = 0
53
- self.use_bad_words[req_idx] = False
47
return
48
49
num_bad_words = len(bad_words_token_ids)
@@ -77,7 +70,6 @@ class BadWordsState:
70
self.bad_word_token_ids.stage_write(req_idx, 0, flattened_tokens)
71
self.bad_word_offsets.stage_write(req_idx, 0, offsets)
72
self.num_bad_words.np[req_idx] = num_bad_words
80
- self.use_bad_words[req_idx] = True
73
74
def apply_staged_writes(self) -> None:
75
self.num_bad_words.copy_to_uva()
@@ -92,11 +84,11 @@ class BadWordsState:
84
input_ids: torch.Tensor,
85
expanded_local_pos: torch.Tensor,
86
) -> None:
95
- if not np.any(self.use_bad_words[idx_mapping_np]):
87
+ max_num_bad_words = int(self.num_bad_words.np[idx_mapping_np].max())
88
+ if max_num_bad_words == 0:
89
# No request uses bad words. Skip the kernel launch.
90
return
91
99
- actual_max_num_bad_words = int(np.max(self.num_bad_words.np[idx_mapping_np]))
92
apply_bad_words(
93
logits,
94
idx_mapping,
@@ -108,7 +100,7 @@ class BadWordsState:
100
self.total_len,
101
input_ids,
102
expanded_local_pos,
111
- actual_max_num_bad_words,
103
+ max_num_bad_words,
104
)
105
106