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

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 15, 22:15 UTC

[v1] Support allowed_token_ids in v1 Sampler (#13210)

Signed-off-by: Lu Fang <lufang@fb.com>

Lu Fang committed Feb 21, 2025 at 22:13 UTC bb78fb318e69f2e5e42ad2f6cf7dd050330c8643

7 files changed +168 -19

tests/v1/sample/test_rejection_sampler.py

@@ -43,6 +43,7 @@ def create_sampling_metadata(spec_tokens: List[List[int]]) -> SamplingMetadata:

43 output_token_ids=[],

44 min_tokens={},

45 logit_bias=[None] * batch_size,

46 + allowed_token_ids_mask=None,

47 )

tests/v1/sample/test_sampler.py

+79 -15

@@ -57,6 +57,26 @@ def _create_logit_bias(

57 return res

60 +def _create_allowed_token_ids(

61 + batch_size: int,

62 + vocab_size: int,

63 + num_allowed_token_ids: int,

64 + device: torch.device,

65 +) -> Optional[torch.Tensor]:

66 + mask: Optional[torch.Tensor] = None

67 + for i in range(batch_size):

68 + if i % 2 == 1:

69 + continue

70 + if mask is None:

71 + mask = torch.zeros((batch_size, vocab_size),

72 + dtype=torch.bool,

73 + device=device)

74 + start = min(i, vocab_size - 1)

75 + end = min(i + num_allowed_token_ids, vocab_size - 1)

76 + mask[i, start:end] = True

77 + return mask

78 +

79 +

80 def _create_default_sampling_metadata(

81 num_output_tokens: int,

82 batch_size: int,

@@ -92,6 +112,7 @@ def _create_default_sampling_metadata(

112 no_penalties=True,

113 min_tokens={},

114 logit_bias=[None] * batch_size,

115 + allowed_token_ids_mask=None,

116 )

117 return fake_sampling_metadata

118

@@ -253,7 +274,10 @@ def test_sampler_frequency_penalty(device: str, batch_size: int,

274 sampling_metadata.frequency_penalties = _create_penalty_tensor(

275 batch_size, frequency_penalty, torch.device(device))

276 output_token_ids, sorted_token_ids_in_output = \

256 - _create_weighted_output_token_list(batch_size, VOCAB_SIZE)

277 + _create_weighted_output_token_list(

278 + batch_size,

279 + VOCAB_SIZE,

280 + )

281 sampling_metadata.output_token_ids = output_token_ids

282 sampling_metadata.no_penalties = False

283 sampler = Sampler()

@@ -262,8 +286,8 @@ def test_sampler_frequency_penalty(device: str, batch_size: int,

286 for batch_idx in range(batch_size):

287 non_penalized_token_id = logits[batch_idx].argmax().item()

288 penalized_token_id = logits[batch_idx].argmin().item()

265 - distinct_sorted_token_ids_in_output = \

266 - sorted_token_ids_in_output[batch_idx]

289 + distinct_sorted_token_ids_in_output = sorted_token_ids_in_output[

290 + batch_idx]

291 most_frequent_token_id = distinct_sorted_token_ids_in_output[

292 len(distinct_sorted_token_ids_in_output) - 1]

293 if frequency_penalty > 0:

@@ -272,8 +296,8 @@ def test_sampler_frequency_penalty(device: str, batch_size: int,

296 # non-penalized token ID is not present in the output, while the

297 # most penalized token is the one that occurs most frequently in

298 # the output.

275 - assert non_penalized_token_id \

276 - not in distinct_sorted_token_ids_in_output

299 + assert (non_penalized_token_id

300 + not in distinct_sorted_token_ids_in_output)

301 assert penalized_token_id == most_frequent_token_id

302 elif frequency_penalty < 0:

303 # If `frequency_penalty` is set to < 0, it indicates

@@ -282,8 +306,7 @@ def test_sampler_frequency_penalty(device: str, batch_size: int,

306 # in the output, while the penalized token ID is one that has not

307 # yet appeared.

308 assert non_penalized_token_id == most_frequent_token_id

285 - assert penalized_token_id \

286 - not in distinct_sorted_token_ids_in_output

309 + assert penalized_token_id not in distinct_sorted_token_ids_in_output

310

311

312 @pytest.mark.parametrize("device", CUDA_DEVICES)

@@ -318,18 +341,18 @@ def test_sampler_repetition_penalty(device: str, batch_size: int,

341 # If `repetition_penalty` > 1.0, verify that the non-penalized

342 # token ID has not been seen before, while the penalized token ID

343 # exists either in the prompt or the output.

321 - assert (non_penalized_token_id not in prompt_tokens and \

322 - non_penalized_token_id not in output_tokens)

323 - assert (penalized_token_id in prompt_tokens or \

324 - penalized_token_id in output_tokens)

344 + assert (non_penalized_token_id not in prompt_tokens

345 + and non_penalized_token_id not in output_tokens)

346 + assert (penalized_token_id in prompt_tokens

347 + or penalized_token_id in output_tokens)

348 elif repetition_penalty < 1.0:

349 # If `repetition_penalty` < 1.0, verify that the penalized

350 # token ID has not been seen before, while the non-penalized

351 # token ID exists either in the prompt or the output.

329 - assert (penalized_token_id not in prompt_tokens and \

330 - penalized_token_id not in output_tokens)

331 - assert (non_penalized_token_id in prompt_tokens or \

332 - non_penalized_token_id in output_tokens)

352 + assert (penalized_token_id not in prompt_tokens

353 + and penalized_token_id not in output_tokens)

354 + assert (non_penalized_token_id in prompt_tokens

355 + or non_penalized_token_id in output_tokens)

356

357

358 @pytest.mark.parametrize("device", CUDA_DEVICES)

@@ -404,3 +427,44 @@ def test_sampler_logit_bias(device: str, batch_size: int, bias_value: float):

427 1e-2)

428 else:

429 assert logits_for_req[token_id] == pytest.approx(1e-2)

430 +

431 +

432 +@pytest.mark.parametrize("device", CUDA_DEVICES)

433 +@pytest.mark.parametrize("batch_size", [1, 2, 32])

434 +@pytest.mark.parametrize("num_allowed_token_ids", [0, 1, 2])

435 +def test_sampler_allowed_token_ids(device: str, batch_size: int,

436 + num_allowed_token_ids: int):

437 + """

438 + Test to verify that when the repetition penalty is enabled, tokens

439 + are penalized based on their presence in the prompt or the existing

440 + output.

441 + """

442 + torch.set_default_device(device)

443 + # Create fake logits where each token is assigned the same

444 + # logit value.

445 + fake_logits = _create_fake_logits(batch_size, VOCAB_SIZE)

446 + sampling_metadata = _create_default_sampling_metadata(

447 + NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE, torch.device(device))

448 + mask = _create_allowed_token_ids(

449 + batch_size=batch_size,

450 + vocab_size=VOCAB_SIZE,

451 + num_allowed_token_ids=num_allowed_token_ids,

452 + device=device,

453 + )

454 + sampling_metadata.allowed_token_ids_mask = mask

455 + sampler = Sampler()

456 + logits = sampler.apply_allowed_token_ids(fake_logits, sampling_metadata)

457 + logits = logits.cpu()

458 + for batch_idx in range(batch_size):

459 + logits_for_req = logits[batch_idx]

460 + if batch_idx % 2 == 1:

461 + assert torch.all(logits_for_req != -float("inf"))

462 + continue

463 + for token_id in range(VOCAB_SIZE):

464 + start = min(batch_idx, VOCAB_SIZE - 1)

465 + end = min(batch_idx + num_allowed_token_ids, VOCAB_SIZE - 1)

466 + if token_id >= start and token_id < end:

467 + assert logits_for_req[token_id] == -float(

468 + "inf"), f"{batch_idx}, {token_id}"

469 + else:

470 + assert logits_for_req[token_id] != -float("inf")

tests/v1/worker/test_gpu_input_batch.py

+13

@@ -66,6 +66,10 @@ def _construct_expected_sampling_metadata(

66 temperature = [0.0 for _ in range(num_reqs)]

67 min_tokens = {}

68 logit_bias = [None] * num_reqs

69 + allowed_token_ids_mask = torch.zeros(num_reqs,

70 + VOCAB_SIZE,

71 + dtype=torch.bool,

72 + device=device)

73 for req in reqs:

74 if req.req_id not in req_ids_retained:

75 continue

@@ -86,6 +90,10 @@ def _construct_expected_sampling_metadata(

90 req.sampling_params.min_tokens,

91 req.sampling_params.all_stop_token_ids)

92 logit_bias[index_in_input_batch] = req.sampling_params.logit_bias

93 + if req.sampling_params.allowed_token_ids:

94 + allowed_token_ids_mask[index_in_input_batch][

95 + req.sampling_params.allowed_token_ids] = True

96 +

97 return SamplingMetadata(

98 temperature=torch.tensor(temperature, dtype=torch.float,

99 device=device),

@@ -121,6 +129,7 @@ def _construct_expected_sampling_metadata(

129 and all(x == 0 for x in frequency_penalties)

130 and all(x == 1 for x in repetition_penalties)),

131 logit_bias=logit_bias,

132 + allowed_token_ids_mask=allowed_token_ids_mask,

133 )

134

135

@@ -242,3 +251,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):

251 assert expected_sampling_metadata.no_penalties == \

252 sampling_metadata.no_penalties

253 assert expected_sampling_metadata.logit_bias == sampling_metadata.logit_bias

254 + if sampling_metadata.allowed_token_ids_mask:

255 + assert torch.allclose(

256 + expected_sampling_metadata.allowed_token_ids_mask,

257 + sampling_metadata.allowed_token_ids_mask)

vllm/v1/engine/processor.py

+14

@@ -83,6 +83,19 @@ class Processor:

83 raise ValueError(f"Got lora_request {lora_request} but LoRA is "

84 "not enabled!")

86 + def _validate_allowed_token_ids(

87 + self,

88 + params: Union[SamplingParams, PoolingParams],

89 + ) -> None:

90 + if not isinstance(params, SamplingParams):

91 + return

92 + if params.allowed_token_ids is None:

93 + return

94 + if not all(0 <= tid < self.model_config.vocab_size

95 + for tid in params.allowed_token_ids):

96 + raise ValueError(

97 + "allowed_token_ids contains out-of-vocab token id")

98 +

99 def process_inputs(

100 self,

101 request_id: str,

@@ -100,6 +113,7 @@ class Processor:

113

114 self._validate_logprobs(params)

115 self._validate_lora(lora_request)

116 + self._validate_allowed_token_ids(params)

117

118 if arrival_time is None:

119 arrival_time = time.time()

vllm/v1/sample/metadata.py

@@ -37,3 +37,7 @@ class SamplingMetadata:

37 min_tokens: Dict[int, Tuple[int, Set[int]]]

39 logit_bias: List[Optional[Dict[int, float]]]

40 +

41 + # `allowed_token_ids_mask` is a 2D bool tensor of shape (max batch size,

42 + # vocab size).

43 + allowed_token_ids_mask: Optional[torch.Tensor]

vllm/v1/sample/sampler.py

+16 -2

@@ -47,6 +47,8 @@ class Sampler(nn.Module):

48 # Use float32 for the logits.

49 logits = logits.to(torch.float32)

50 + # Apply allowed token ids.

51 + logits = self.apply_allowed_token_ids(logits, sampling_metadata)

52 # Apply logits bias.

53 logits = self.apply_logits_bias(logits, sampling_metadata)

54 # Apply penalties (e.g., min_tokens, freq_penalties).

@@ -184,11 +186,13 @@ class Sampler(nn.Module):

186 if not sampling_metadata.no_penalties:

187 assert sampling_metadata.prompt_token_ids is not None

188 logits = apply_all_penalties(

187 - logits, sampling_metadata.prompt_token_ids,

189 + logits,

190 + sampling_metadata.prompt_token_ids,

191 sampling_metadata.presence_penalties,

192 sampling_metadata.frequency_penalties,

193 sampling_metadata.repetition_penalties,

191 - sampling_metadata.output_token_ids)

194 + sampling_metadata.output_token_ids,

195 + )

196 return logits

197

198 def apply_min_p(

@@ -226,3 +230,13 @@ class Sampler(nn.Module):

230 for token_id, bias in logit_bias.items():

231 logits[i, token_id] += bias

232 return logits

233 +

234 + def apply_allowed_token_ids(

235 + self,

236 + logits: torch.Tensor,

237 + sampling_metadata: SamplingMetadata,

238 + ) -> torch.Tensor:

239 + if sampling_metadata.allowed_token_ids_mask is not None:

240 + logits.masked_fill_(sampling_metadata.allowed_token_ids_mask,

241 + float("-inf"))

242 + return logits

vllm/v1/worker/gpu_input_batch.py

+41 -2

@@ -143,7 +143,7 @@ class InputBatch:

143 device="cpu",

144 pin_memory=pin_memory)

145 self.frequency_penalties_cpu = \

146 - self.frequency_penalties_cpu_tensor.numpy()

146 + self.frequency_penalties_cpu_tensor.numpy()

147 self.frequency_penalties_reqs: Set[str] = set()

148

149 # Presence penalty related data structures

@@ -168,7 +168,7 @@ class InputBatch:

168 device="cpu",

169 pin_memory=pin_memory)

170 self.repetition_penalties_cpu = \

171 - self.repetition_penalties_cpu_tensor.numpy()

171 + self.repetition_penalties_cpu_tensor.numpy()

172 self.repetition_penalties_reqs: Set[str] = set()

173

174 # req_index -> (min_tokens, stop_token_ids)

@@ -192,6 +192,9 @@ class InputBatch:

192

193 self.logit_bias: List[Optional[Dict[int,

194 float]]] = [None] * max_num_reqs

195 + self.has_allowed_token_ids: Set[str] = set()

196 + self.allowed_token_ids_mask: Optional[torch.Tensor] = None

197 + self.allowed_token_ids_mask_cpu_tensor: Optional[torch.Tensor] = None

198

199 self.req_output_token_ids: List[Optional[List[int]]] = []

200

@@ -287,6 +290,22 @@ class InputBatch:

290 if sampling_params.logit_bias is not None:

291 self.logit_bias[req_index] = sampling_params.logit_bias

292

293 + if sampling_params.allowed_token_ids:

294 + self.has_allowed_token_ids.add(req_id)

295 + if self.allowed_token_ids_mask_cpu_tensor is None:

296 + # Lazy allocation for this tensor, which can be large.

297 + self.allowed_token_ids_mask = torch.zeros(self.max_num_reqs,

298 + self.vocab_size,

299 + dtype=torch.bool,

300 + device=self.device)

301 + self.allowed_token_ids_mask_cpu_tensor = torch.zeros(

302 + self.max_num_reqs,

303 + self.vocab_size,

304 + dtype=torch.bool,

305 + device="cpu")

306 + self.allowed_token_ids_mask_cpu_tensor[req_index][

307 + sampling_params.allowed_token_ids] = True

308 +

309 # Add request lora ID

310 if request.lora_request:

311 lora_id = request.lora_request.lora_int_id

@@ -332,6 +351,9 @@ class InputBatch:

351 self.request_lora_mapping[req_index] = 0

352

353 self.logit_bias[req_index] = None

354 + self.has_allowed_token_ids.discard(req_id)

355 + if self.allowed_token_ids_mask_cpu_tensor is not None:

356 + self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)

357 return req_index

358

359 def condense(self, empty_req_indices: List[int]) -> None:

@@ -400,6 +422,11 @@ class InputBatch:

422

423 self.logit_bias[empty_index] = self.logit_bias[last_req_index]

424

425 + if self.allowed_token_ids_mask_cpu_tensor is not None:

426 + self.allowed_token_ids_mask_cpu_tensor[

427 + empty_index] = self.allowed_token_ids_mask_cpu_tensor[

428 + last_req_index]

429 +

430 # Decrement last_req_index since it is now empty.

431 last_req_index -= 1

432

@@ -442,6 +469,13 @@ class InputBatch:

469 else:

470 prompt_token_ids = None

471

472 + allowed_token_ids_mask: Optional[torch.Tensor] = None

473 + if not self.no_allowed_token_ids:

474 + assert self.allowed_token_ids_mask is not None

475 + copy_slice(self.allowed_token_ids_mask_cpu_tensor,

476 + self.allowed_token_ids_mask, num_reqs)

477 + allowed_token_ids_mask = self.allowed_token_ids_mask[:num_reqs]

478 +

479 return SamplingMetadata(

480 temperature=temperature,

481 all_greedy=self.all_greedy,

@@ -460,6 +494,7 @@ class InputBatch:

494 min_tokens=self.min_tokens,

495 no_penalties=self.no_penalties,

496 logit_bias=self.logit_bias[:num_reqs],

497 + allowed_token_ids_mask=allowed_token_ids_mask,

498 )

499

500 def get_sampling_metadata(

@@ -550,3 +585,7 @@ class InputBatch:

585 @property

586 def no_prompt_logprob(self) -> bool:

587 return not self.num_prompt_logprobs

588 +

589 + @property

590 + def no_allowed_token_ids(self) -> bool:

591 + return len(self.has_allowed_token_ids) == 0