@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 387a1898d9

Fork 0

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] support bad_words sampling param (#33433)

Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>

zhrrr committed Feb 17, 2026 at 08:36 UTC 387a1898d9593f001734527946af9aafa4e24ae6

7 files changed +303 -43

vllm/v1/worker/gpu/input_batch.py

+24 -8

@@ -156,8 +156,8 @@ def _prepare_prefill_inputs_kernel(

156 next_prefill_tokens_ptr,

157 idx_mapping_ptr,

158 query_start_loc_ptr,

159 - prefill_token_ids_ptr,

160 - prefill_token_ids_stride,

159 + all_token_ids_ptr,

160 + all_token_ids_stride,

161 prefill_lens_ptr,

162 num_computed_tokens_ptr,

163 BLOCK_SIZE: tl.constexpr,

@@ -174,16 +174,16 @@ def _prepare_prefill_inputs_kernel(

174 query_end = tl.load(query_start_loc_ptr + batch_idx + 1)

175 query_len = query_end - query_start

176

177 - prefill_ptr = prefill_token_ids_ptr + req_state_idx * prefill_token_ids_stride

177 + request_ptr = all_token_ids_ptr + req_state_idx * all_token_ids_stride

178 for i in range(0, query_len, BLOCK_SIZE):

179 block = i + tl.arange(0, BLOCK_SIZE)

180 mask = block < query_len

181 - tokens = tl.load(prefill_ptr + num_computed + block, mask=mask)

181 + tokens = tl.load(request_ptr + num_computed + block, mask=mask)

182 tl.store(input_ids_ptr + query_start + block, tokens, mask=mask)

183

184 next_pos = num_computed + query_len

185 if next_pos < prefill_len:

186 - next_token = tl.load(prefill_ptr + next_pos)

186 + next_token = tl.load(request_ptr + next_pos)

187 tl.store(next_prefill_tokens_ptr + req_state_idx, next_token)

188

189

@@ -192,7 +192,7 @@ def prepare_prefill_inputs(

192 next_prefill_tokens: torch.Tensor,

193 idx_mapping: torch.Tensor,

194 query_start_loc: torch.Tensor,

195 - prefill_token_ids: torch.Tensor,

195 + all_token_ids: torch.Tensor,

196 prefill_len: torch.Tensor,

197 num_computed_tokens: torch.Tensor,

198 ) -> None:

@@ -202,8 +202,8 @@ def prepare_prefill_inputs(

202 next_prefill_tokens,

203 idx_mapping,

204 query_start_loc,

205 - prefill_token_ids,

206 - prefill_token_ids.stride(0),

205 + all_token_ids,

206 + all_token_ids.stride(0),

207 prefill_len,

208 num_computed_tokens,

209 BLOCK_SIZE=1024,

@@ -423,16 +423,21 @@ def _post_update_kernel(

423 num_sampled_ptr,

424 num_rejected_ptr,

425 query_start_loc_ptr,

426 + all_token_ids_ptr,

427 + all_token_ids_stride,

428 + total_len_ptr,

429 ):

430 req_id = tl.program_id(0)

431 req_state_idx = tl.load(idx_mapping_ptr + req_id)

432

433 + total_len = tl.load(total_len_ptr + req_state_idx)

434 num_sampled = tl.load(num_sampled_ptr + req_id)

435 if num_sampled > 0:

436 token_id = tl.load(

437 sampled_tokens_ptr + req_id * sampled_tokens_stride + num_sampled - 1

438 )

439 tl.store(last_sampled_tokens_ptr + req_state_idx, token_id)

440 + tl.store(total_len_ptr + req_state_idx, total_len + num_sampled)

441

442 for i in range(num_sampled):

443 token_id = tl.load(sampled_tokens_ptr + req_id * sampled_tokens_stride + i)

@@ -442,6 +447,10 @@ def _post_update_kernel(

447 count = tl.load(token_ptr)

448 count += 1

449 tl.store(token_ptr, count)

450 + tl.store(

451 + all_token_ids_ptr + req_state_idx * all_token_ids_stride + total_len + i,

452 + token_id,

453 + )

454

455 query_start = tl.load(query_start_loc_ptr + req_id)

456 query_end = tl.load(query_start_loc_ptr + req_id + 1)

@@ -470,6 +479,10 @@ def post_update(

479 num_rejected: torch.Tensor,

480 # [num_reqs + 1]

481 query_start_loc: torch.Tensor,

482 + # [max_num_reqs, max_model_len]

483 + all_token_ids: torch.Tensor,

484 + # [max_num_reqs]

485 + total_len: torch.Tensor,

486 ) -> None:

487 num_reqs = idx_mapping.shape[0]

488 _post_update_kernel[(num_reqs,)](

@@ -483,6 +496,9 @@ def post_update(

496 num_sampled,

497 num_rejected,

498 query_start_loc,

499 + all_token_ids,

500 + all_token_ids.stride(0),

501 + total_len,

502 num_warps=1,

503 )

504

vllm/v1/worker/gpu/model_runner.py

+11 -6

@@ -151,6 +151,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

151 max_num_reqs=self.max_num_reqs,

152 vocab_size=self.vocab_size,

153 device=self.device,

154 + all_token_ids=self.req_states.all_token_ids.gpu,

155 + prompt_len=self.req_states.prompt_len.gpu,

156 + total_len=self.req_states.total_len.gpu,

157 logprobs_mode=self.model_config.logprobs_mode,

158 num_speculative_tokens=self.num_speculative_steps + 1,

159 )

@@ -448,7 +451,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

451 self.req_states.add_request(

452 req_id=req_id,

453 prompt_len=prompt_len,

451 - prefill_token_ids=new_req_data.prefill_token_ids,

454 + all_token_ids=new_req_data.prefill_token_ids,

455 num_computed_tokens=new_req_data.num_computed_tokens,

456 )

457 req_index = self.req_states.req_id_to_index[req_id]

@@ -479,9 +482,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

482 if scheduler_output.scheduled_new_reqs:

483 self.req_states.apply_staged_writes()

484 self.sampler.apply_staged_writes(

482 - self.req_states.prefill_token_ids.gpu,

485 + self.req_states.all_token_ids.gpu,

486 self.req_states.prefill_len.np,

484 - self.req_states.prompt_len,

487 + self.req_states.prompt_len.np,

488 )

489 if self.uses_mrope:

490 self.mrope_states.apply_staged_writes()

@@ -570,7 +573,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

573 self.req_states.next_prefill_tokens,

574 idx_mapping,

575 query_start_loc,

573 - self.req_states.prefill_token_ids.gpu,

576 + self.req_states.all_token_ids.gpu,

577 self.req_states.prefill_len.gpu,

578 self.req_states.num_computed_tokens.gpu,

579 )

@@ -759,6 +762,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

762 num_sampled,

763 num_rejected,

764 input_batch.query_start_loc,

765 + self.req_states.all_token_ids.gpu,

766 + self.req_states.total_len.gpu,

767 )

768

769 # Update the number of computed prefill tokens.

@@ -924,9 +929,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

929 self.model.compute_logits,

930 hidden_states,

931 input_batch,

927 - self.req_states.prefill_token_ids.gpu,

932 + self.req_states.all_token_ids.gpu,

933 self.req_states.num_computed_tokens.gpu,

929 - self.req_states.prompt_len,

934 + self.req_states.prompt_len.np,

935 self.req_states.prefill_len.np,

936 self.req_states.num_computed_prefill_tokens,

937 )

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

+209

@@ -0,0 +1,209 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import numpy as np

4 +import torch

5 +

6 +from vllm.sampling_params import SamplingParams

7 +from vllm.triton_utils import tl, triton

8 +from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor

9 +

10 +MAX_BAD_WORDS_TOTAL_TOKENS = 1024 # Max total tokens for all bad words per request

11 +MAX_NUM_BAD_WORDS = 128 # Max number of bad words per request

12 +

13 +

14 +class BadWordsState:

15 + def __init__(

16 + self,

17 + all_token_ids: torch.Tensor,

18 + prompt_len: torch.Tensor,

19 + total_len: torch.Tensor,

20 + ):

21 + self.all_token_ids = all_token_ids

22 + self.prompt_len = prompt_len

23 + self.total_len = total_len

24 +

25 + self.max_num_reqs = prompt_len.shape[0]

26 + self.device = prompt_len.device

27 +

28 + # flattened bad word tokens: [max_num_reqs, MAX_BAD_WORDS_TOTAL_TOKENS]

29 + self.bad_word_token_ids = StagedWriteTensor(

30 + (self.max_num_reqs, MAX_BAD_WORDS_TOTAL_TOKENS),

31 + dtype=torch.int32,

32 + device=self.device,

33 + )

34 + # cumulative offsets of bad words: [max_num_reqs, MAX_NUM_BAD_WORDS + 1]

35 + self.bad_word_offsets = StagedWriteTensor(

36 + (self.max_num_reqs, MAX_NUM_BAD_WORDS + 1),

37 + dtype=torch.int32,

38 + device=self.device,

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)

44 +

45 + def add_request(

46 + self,

47 + req_idx: int,

48 + sampling_params: SamplingParams,

49 + ) -> None:

50 + bad_words_token_ids = sampling_params.bad_words_token_ids

51 + if not bad_words_token_ids:

52 + self.num_bad_words.np[req_idx] = 0

53 + self.use_bad_words[req_idx] = False

54 + return

55 +

56 + num_bad_words = len(bad_words_token_ids)

57 + if num_bad_words > MAX_NUM_BAD_WORDS:

58 + raise ValueError(

59 + f"Too many bad words: {num_bad_words}. "

60 + f"The max number is {MAX_NUM_BAD_WORDS}."

61 + )

62 +

63 + # Flatten bad words and compute offsets

64 + flattened_tokens: list[int] = []

65 + offsets: list[int] = [0]

66 + for bad_word in bad_words_token_ids:

67 + flattened_tokens.extend(bad_word)

68 + offsets.append(len(flattened_tokens))

69 +

70 + if len(flattened_tokens) > MAX_BAD_WORDS_TOTAL_TOKENS:

71 + raise ValueError(

72 + f"Too many total bad word tokens: {len(flattened_tokens)}. "

73 + f"The max is {MAX_BAD_WORDS_TOTAL_TOKENS}."

74 + )

75 +

76 + # Stage writes

77 + self.bad_word_token_ids.stage_write(req_idx, 0, flattened_tokens)

78 + self.bad_word_offsets.stage_write(req_idx, 0, offsets)

79 + self.num_bad_words.np[req_idx] = num_bad_words

80 + self.use_bad_words[req_idx] = True

81 +

82 + def apply_staged_writes(self) -> None:

83 + self.num_bad_words.copy_to_uva()

84 + self.bad_word_token_ids.apply_write()

85 + self.bad_word_offsets.apply_write()

86 +

87 + def apply_bad_words(

88 + self,

89 + logits: torch.Tensor,

90 + idx_mapping: torch.Tensor,

91 + idx_mapping_np: np.ndarray,

92 + input_ids: torch.Tensor,

93 + expanded_local_pos: torch.Tensor,

94 + ) -> None:

95 + if not np.any(self.use_bad_words[idx_mapping_np]):

96 + # No request uses bad words. Skip the kernel launch.

97 + return

98 +

99 + actual_max_num_bad_words = int(np.max(self.num_bad_words.np[idx_mapping_np]))

100 + apply_bad_words(

101 + logits,

102 + idx_mapping,

103 + self.bad_word_token_ids.gpu,

104 + self.bad_word_offsets.gpu,

105 + self.num_bad_words.gpu,

106 + self.all_token_ids,

107 + self.prompt_len,

108 + self.total_len,

109 + input_ids,

110 + expanded_local_pos,

111 + actual_max_num_bad_words,

112 + )

113 +

114 +

115 +@triton.jit

116 +def _bad_words_kernel(

117 + logits_ptr,

118 + logits_stride,

119 + expanded_idx_mapping_ptr,

120 + bad_word_token_ids_ptr,

121 + bad_word_token_ids_stride,

122 + bad_word_offsets_ptr,

123 + bad_word_offsets_stride,

124 + num_bad_words_ptr,

125 + all_token_ids_ptr,

126 + all_token_ids_stride,

127 + prompt_len_ptr,

128 + total_len_ptr,

129 + input_ids_ptr,

130 + expanded_local_pos_ptr,

131 +):

132 + logit_idx = tl.program_id(0)

133 + bw_idx = tl.program_id(1)

134 +

135 + req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx)

136 + num_bad_words = tl.load(num_bad_words_ptr + req_state_idx)

137 +

138 + if bw_idx >= num_bad_words:

139 + return

140 +

141 + pos = tl.load(expanded_local_pos_ptr + logit_idx)

142 + cur_req_first_pos = logit_idx - pos

143 +

144 + prompt_len = tl.load(prompt_len_ptr + req_state_idx)

145 + total_len = tl.load(total_len_ptr + req_state_idx)

146 + output_len = total_len - prompt_len

147 + effective_len = output_len + pos

148 +

149 + bd_offsets_base = bad_word_offsets_ptr + req_state_idx * bad_word_offsets_stride

150 + bd_tokens_base = bad_word_token_ids_ptr + req_state_idx * bad_word_token_ids_stride

151 + output_base = all_token_ids_ptr + req_state_idx * all_token_ids_stride + prompt_len

152 +

153 + start = tl.load(bd_offsets_base + bw_idx)

154 + end = tl.load(bd_offsets_base + bw_idx + 1)

155 + bad_word_len = end - start

156 + prefix_len = bad_word_len - 1

157 +

158 + if prefix_len > effective_len:

159 + return

160 +

161 + last_token = tl.load(bd_tokens_base + end - 1)

162 + match = 1

163 + for i in range(prefix_len):

164 + expected = tl.load(bd_tokens_base + start + i)

165 + actual_pos = effective_len - prefix_len + i

166 +

167 + from_spec_input = actual_pos >= output_len

168 + if from_spec_input:

169 + spec_offset = actual_pos - output_len

170 + actual = tl.load(input_ids_ptr + cur_req_first_pos + spec_offset)

171 + else:

172 + actual = tl.load(output_base + actual_pos)

173 +

174 + match = match & (expected == actual)

175 +

176 + if match:

177 + tl.store(logits_ptr + logit_idx * logits_stride + last_token, -float("inf"))

178 +

179 +

180 +def apply_bad_words(

181 + logits: torch.Tensor,

182 + expanded_idx_mapping: torch.Tensor,

183 + bad_word_token_ids: torch.Tensor,

184 + bad_word_offsets: torch.Tensor,

185 + num_bad_words: torch.Tensor,

186 + all_token_ids: torch.Tensor,

187 + prompt_len: torch.Tensor,

188 + total_len: torch.Tensor,

189 + input_ids: torch.Tensor,

190 + expanded_local_pos: torch.Tensor,

191 + max_num_bad_words: int,

192 +) -> None:

193 + total_num_tokens = logits.shape[0]

194 + _bad_words_kernel[(total_num_tokens, max_num_bad_words)](

195 + logits,

196 + logits.stride(0),

197 + expanded_idx_mapping,

198 + bad_word_token_ids,

199 + bad_word_token_ids.stride(0),

200 + bad_word_offsets,

201 + bad_word_offsets.stride(0),

202 + num_bad_words,

203 + all_token_ids,

204 + all_token_ids.stride(0),

205 + prompt_len,

206 + total_len,

207 + input_ids,

208 + expanded_local_pos,

209 + )

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

+10 -10

@@ -51,14 +51,14 @@ class PenaltiesState:

52 def apply_staged_writes(

53 self,

54 - prefill_token_ids: torch.Tensor,

54 + all_token_ids: torch.Tensor,

55 prefill_lens: np.ndarray,

56 prompt_lens: np.ndarray,

57 ) -> None:

58 # TODO(woosuk): Optimize this.

59 for req_idx in self._penalties_reqs:

60 bincount(

61 - prefill_token_ids[req_idx],

61 + all_token_ids[req_idx],

62 int(prefill_lens[req_idx]),

63 int(prompt_lens[req_idx]),

64 self.prompt_bin_mask[req_idx],

@@ -216,7 +216,7 @@ def apply_penalties(

216

217 @triton.jit(do_not_specialize=["prefill_len", "prompt_len"])

218 def _bincount_kernel(

219 - prefill_token_ids_ptr,

219 + all_token_ids_ptr,

220 prefill_len,

221 prompt_len,

222 prompt_bin_mask_ptr,

@@ -230,20 +230,20 @@ def _bincount_kernel(

230 block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

231 if block_idx * BLOCK_SIZE < prompt_len:

232 mask = block < prompt_len

233 - prefill_tokens = tl.load(prefill_token_ids_ptr + block, mask=mask)

234 - idx = prefill_tokens // 32

235 - bit_idx = prefill_tokens % 32

233 + prompt_tokens = tl.load(all_token_ids_ptr + block, mask=mask)

234 + idx = prompt_tokens // 32

235 + bit_idx = prompt_tokens % 32

236 bit = tl.full((BLOCK_SIZE,), 1, tl.int32) << bit_idx

237 tl.atomic_or(prompt_bin_mask_ptr + idx, bit, mask=mask)

238 if (block_idx + 1) * BLOCK_SIZE >= prompt_len:

239 mask = block < prefill_len

240 mask &= block >= prompt_len

241 - prefill_tokens = tl.load(prefill_token_ids_ptr + block, mask=mask)

242 - tl.atomic_add(output_bin_counts_ptr + prefill_tokens, 1, mask=mask)

241 + output_tokens = tl.load(all_token_ids_ptr + block, mask=mask)

242 + tl.atomic_add(output_bin_counts_ptr + output_tokens, 1, mask=mask)

243

244

245 def bincount(

246 - prefill_token_ids: torch.Tensor,

246 + all_token_ids: torch.Tensor,

247 prefill_len: int,

248 prompt_len: int,

249 prompt_bin_mask: torch.Tensor,

@@ -254,7 +254,7 @@ def bincount(

254 BLOCK_SIZE = 1024

255 num_blocks = triton.cdiv(prefill_len, BLOCK_SIZE)

256 _bincount_kernel[(num_blocks,)](

257 - prefill_token_ids,

257 + all_token_ids,

258 prefill_len,

259 prompt_len,

260 prompt_bin_mask,

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

+8 -10

@@ -36,7 +36,7 @@ class PromptLogprobsWorker:

36 hidden_states: torch.Tensor,

37 input_batch: InputBatch,

38 # [max_num_reqs, max_model_len]

39 - prefill_token_ids: torch.Tensor,

39 + all_token_ids: torch.Tensor,

40 # [max_num_reqs]

41 num_computed_tokens: torch.Tensor,

42 # [max_num_reqs]

@@ -70,7 +70,7 @@ class PromptLogprobsWorker:

70 input_batch.query_start_loc,

71 input_batch.idx_mapping,

72 num_computed_tokens,

73 - prefill_token_ids,

73 + all_token_ids,

74 )

75 # Compute the prompt logprobs.

76 prompt_logprobs, prompt_ranks = compute_prompt_logprobs_with_chunking(

@@ -132,8 +132,8 @@ def _prompt_logprobs_token_ids_kernel(

132 query_start_loc_ptr,

133 idx_mapping_ptr,

134 num_computed_tokens_ptr,

135 - prefill_token_ids_ptr,

136 - prefill_token_ids_stride,

135 + all_token_ids_ptr,

136 + all_token_ids_stride,

137 BLOCK_SIZE: tl.constexpr,

138 ):

139 batch_idx = tl.program_id(0)

@@ -151,9 +151,7 @@ def _prompt_logprobs_token_ids_kernel(

151 # because the logprob is computed for the next token.

152 target_pos = num_computed_tokens + 1 + block

153 token_ids = tl.load(

154 - prefill_token_ids_ptr

155 - + req_state_idx * prefill_token_ids_stride

156 - + target_pos,

154 + all_token_ids_ptr + req_state_idx * all_token_ids_stride + target_pos,

155 mask=mask,

156 )

157 tl.store(

@@ -166,7 +164,7 @@ def get_prompt_logprobs_token_ids(

164 query_start_loc: torch.Tensor,

165 idx_mapping: torch.Tensor,

166 num_computed_tokens: torch.Tensor,

169 - prefill_token_ids: torch.Tensor,

167 + all_token_ids: torch.Tensor,

168 ) -> torch.Tensor:

169 token_ids = torch.empty(num_tokens, dtype=torch.int64, device=idx_mapping.device)

170 num_reqs = idx_mapping.shape[0]

@@ -175,8 +173,8 @@ def get_prompt_logprobs_token_ids(

173 query_start_loc,

174 idx_mapping,

175 num_computed_tokens,

178 - prefill_token_ids,

179 - prefill_token_ids.stride(0),

176 + all_token_ids,

177 + all_token_ids.stride(0),

178 BLOCK_SIZE=1024,

179 )

180 return token_ids

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

+18 -2

@@ -8,6 +8,7 @@ import vllm.envs as envs

8 from vllm.config.model import LogprobsMode

9 from vllm.sampling_params import SamplingParams

10 from vllm.v1.worker.gpu.metrics.logits import get_num_nans

11 +from vllm.v1.worker.gpu.sample.bad_words import BadWordsState

12 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample

13 from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState

14 from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs

@@ -22,6 +23,9 @@ class Sampler:

23 max_num_reqs: int,

24 vocab_size: int,

25 device: torch.device,

26 + all_token_ids: torch.Tensor,

27 + prompt_len: torch.Tensor,

28 + total_len: torch.Tensor,

29 logprobs_mode: LogprobsMode = "raw_logprobs",

30 num_speculative_tokens: int = 1,

31 ):

@@ -33,6 +37,7 @@ class Sampler:

37 self.sampling_states = SamplingStates(max_num_reqs, vocab_size)

38 self.penalties_state = PenaltiesState(max_num_reqs, vocab_size, device)

39 self.logit_bias_state = LogitBiasState(max_num_reqs, device)

40 + self.bad_words_state = BadWordsState(all_token_ids, prompt_len, total_len)

41 self.num_speculative_tokens = num_speculative_tokens

43 def add_request(

@@ -41,18 +46,20 @@ class Sampler:

46 self.sampling_states.add_request(req_idx, sampling_params)

47 self.penalties_state.add_request(req_idx, sampling_params)

48 self.logit_bias_state.add_request(req_idx, prompt_len, sampling_params)

49 + self.bad_words_state.add_request(req_idx, sampling_params)

51 def apply_staged_writes(

52 self,

47 - prefill_token_ids: torch.Tensor,

53 + all_token_ids: torch.Tensor,

54 prefill_lens: np.ndarray,

55 prompt_lens: np.ndarray,

56 ) -> None:

57 self.sampling_states.apply_staged_writes()

58 self.penalties_state.apply_staged_writes(

53 - prefill_token_ids, prefill_lens, prompt_lens

59 + all_token_ids, prefill_lens, prompt_lens

60 )

61 self.logit_bias_state.apply_staged_writes()

62 + self.bad_words_state.apply_staged_writes()

64 def __call__(

65 self,

@@ -124,6 +131,15 @@ class Sampler:

131 self.num_speculative_tokens,

132 )

133

134 + # Apply bad words masking in place.

135 + self.bad_words_state.apply_bad_words(

136 + logits,

137 + idx_mapping,

138 + idx_mapping_np,

139 + input_ids,

140 + expanded_local_pos,

141 + )

142 +

143 # Apply temperature in place.

144 self.sampling_states.apply_temperature(logits, idx_mapping, idx_mapping_np)

145

vllm/v1/worker/gpu/states.py

+23 -7

@@ -27,17 +27,30 @@ class RequestState:

27 self.index_to_req_id: dict[int, str] = {}

28 self.free_indices = list(range(max_num_reqs))

30 - self.prompt_len = np.zeros(self.max_num_reqs, dtype=np.int32)

30 # NOTE(woosuk): This tensor can be extremely large (e.g., several GBs)

31 # depending on the configured max_num_reqs and max_model_len.

32 # To save GPU memory, we use UVA instead of GPU for this tensor.

34 - self.prefill_token_ids = StagedWriteTensor(

33 + self.all_token_ids = StagedWriteTensor(

34 (self.max_num_reqs, self.max_model_len),

35 dtype=torch.int32,

36 device=device,

37 uva_instead_of_gpu=True,

38 )

39 + # NOTE(woosuk): Distinguish clearly between prompt_len and prefill_len:

40 + # - prompt_len: Number of tokens in the user-provided prompt.

41 + # - prefill_len: Number of tokens passed into the model runner.

42 + # This can include the prompt and additional partial output tokens,

43 + # so prefill_len >= prompt_len.

44 + # Usually, prefill_len equals prompt_len, but in cases such as resumption after

45 + # preemption, prefill_len may be greater. Differentiating between these values

46 + # is crucial, as certain features such as prompt logprobs or frequency penalties

47 + # must treat prompt and output tokens separately.

48 + self.prompt_len = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)

49 self.prefill_len = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)

50 + # total_len = prompt_len + output_len. It grows as the request progresses.

51 + self.total_len = StagedWriteTensor(

52 + self.max_num_reqs, dtype=torch.int32, device=device

53 + )

55 # Number of computed tokens.

56 self.num_computed_prefill_tokens = np.zeros(self.max_num_reqs, dtype=np.int32)

@@ -72,7 +85,7 @@ class RequestState:

85 self,

86 req_id: str,

87 prompt_len: int,

75 - prefill_token_ids: list[int],

88 + all_token_ids: list[int],

89 num_computed_tokens: int,

90 ) -> None:

91 assert len(self.free_indices) > 0, "No free indices"

@@ -80,19 +93,22 @@ class RequestState:

93 self.req_id_to_index[req_id] = req_idx

94 self.index_to_req_id[req_idx] = req_id

83 - self.prompt_len[req_idx] = prompt_len

84 - prefill_len = len(prefill_token_ids)

96 + self.prompt_len.np[req_idx] = prompt_len

97 + prefill_len = len(all_token_ids)

98 assert prefill_len >= prompt_len, (

99 f"prefill_len {prefill_len} < prompt_len {prompt_len}"

100 )

101 self.prefill_len.np[req_idx] = prefill_len

89 - self.prefill_token_ids.stage_write(req_idx, 0, prefill_token_ids)

102 + self.total_len.stage_write_elem(req_idx, prefill_len)

103 + self.all_token_ids.stage_write(req_idx, 0, all_token_ids)

104 self.num_computed_prefill_tokens[req_idx] = num_computed_tokens

105 self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens)

106

107 def apply_staged_writes(self) -> None:

108 + self.prompt_len.copy_to_uva()

109 self.prefill_len.copy_to_uva()

95 - self.prefill_token_ids.apply_write()

110 + self.total_len.apply_write()

111 + self.all_token_ids.apply_write()

112 self.num_computed_tokens.apply_write()

113

114 def remove_request(self, req_id: str) -> None: