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

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 14, 05:30 UTC

[V0 Deprecation] Remove V0 Samplers test (#23862)

Woosuk Kwon committed Aug 28, 2025 at 18:05 UTC b668055a114086b8968d9ff4a53586f1d8ea0b47

2 files changed -855

tests/samplers/test_sampler.py deleted

-769

@@ -1,769 +0,0 @@

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

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

3 -

4 -import itertools

5 -import random

6 -from dataclasses import dataclass

7 -from typing import Optional

8 -from unittest.mock import Mock, patch

9 -

10 -import pytest

11 -import torch

12 -from transformers import GenerationConfig, GenerationMixin

13 -

14 -import vllm.envs as envs

15 -from vllm.model_executor.layers.sampler import Sampler

16 -from vllm.model_executor.sampling_metadata import SamplingMetadata

17 -from vllm.model_executor.utils import set_random_seed

18 -from vllm.sequence import SamplingParams, SequenceData, SequenceGroupMetadata

19 -from vllm.utils import Counter, is_pin_memory_available

20 -

21 -

22 -@pytest.fixture(scope="function", autouse=True)

23 -def use_v0_only(monkeypatch):

24 - """

25 - This file tests V0 internals, so set VLLM_USE_V1=0.

26 - """

27 - monkeypatch.setenv('VLLM_USE_V1', '0')

28 -

29 -

30 -class MockLogitsSampler(Sampler):

31 -

32 - def __init__(self, fake_logits: torch.Tensor):

33 - super().__init__()

34 - self.fake_logits = fake_logits

35 -

36 - def forward(self, *args, **kwargs):

37 - return super().forward(*args, **kwargs)

38 -

39 -

40 -def _prepare_test(

41 - batch_size: int

42 -) -> tuple[torch.Tensor, torch.Tensor, MockLogitsSampler]:

43 - input_tensor = torch.rand((batch_size, 1024), dtype=torch.float16)

44 - fake_logits = torch.full((batch_size, VOCAB_SIZE),

45 - 1e-2,

46 - dtype=input_tensor.dtype)

47 - sampler = MockLogitsSampler(fake_logits)

48 - return input_tensor, fake_logits, sampler

49 -

50 -

51 -VOCAB_SIZE = 32000

52 -RANDOM_SEEDS = list(range(128))

53 -CUDA_DEVICES = [

54 - f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)

55 -]

56 -

57 -

58 -def _do_sample(

59 - batch_size: int,

60 - input_tensor: torch.Tensor,

61 - sampler: MockLogitsSampler,

62 - sampling_params: SamplingParams,

63 - device: str,

64 -):

65 - seq_group_metadata_list: list[SequenceGroupMetadata] = []

66 - seq_lens: list[int] = []

67 - for i in range(batch_size):

68 - seq_group_metadata_list.append(

69 - SequenceGroupMetadata(

70 - request_id=f"test_{i}",

71 - is_prompt=True,

72 - seq_data={0: SequenceData.from_seqs([1, 2, 3])},

73 - sampling_params=sampling_params,

74 - block_tables={0: [1]},

75 - ))

76 - seq_lens.append(seq_group_metadata_list[-1].seq_data[0].get_len())

77 -

78 - sampling_metadata = SamplingMetadata.prepare(

79 - seq_group_metadata_list,

80 - seq_lens,

81 - query_lens=seq_lens,

82 - device=device,

83 - pin_memory=is_pin_memory_available())

84 - return sampler(logits=input_tensor, sampling_metadata=sampling_metadata)

85 -

86 -

87 -@pytest.mark.parametrize("seed", RANDOM_SEEDS)

88 -@pytest.mark.parametrize("device", CUDA_DEVICES)

89 -def test_sampler_all_greedy(seed: int, device: str):

90 - set_random_seed(seed)

91 - torch.set_default_device(device)

92 - batch_size = random.randint(1, 256)

93 - input_tensor, fake_logits, sampler = _prepare_test(batch_size)

94 -

95 - sampling_params = SamplingParams(temperature=0)

96 - sampler_output = _do_sample(batch_size, fake_logits, sampler,

97 - sampling_params, device)

98 - expected = torch.argmax(fake_logits, dim=-1)

99 - for i, sequence_output in enumerate(sampler_output):

100 - for nth_output in sequence_output.samples:

101 - assert nth_output.output_token == expected[i].item()

102 -

103 -

104 -@pytest.mark.parametrize("seed", RANDOM_SEEDS)

105 -@pytest.mark.parametrize("device", CUDA_DEVICES)

106 -def test_sampler_all_random(seed: int, device: str):

107 - set_random_seed(seed)

108 - torch.set_default_device(device)

109 - batch_size = random.randint(1, 256)

110 - _, fake_logits, sampler = _prepare_test(batch_size)

111 -

112 - for i in range(batch_size):

113 - fake_logits[i, i] = 1e2

114 -

115 - sampling_params = SamplingParams(

116 - temperature=1.0,

117 - n=random.randint(1, 10),

118 - )

119 - sampler_output = _do_sample(batch_size, fake_logits, sampler,

120 - sampling_params, device)

121 -

122 - for i, sequence_output in enumerate(sampler_output):

123 - for nth_output in sequence_output.samples:

124 - assert nth_output.output_token == i

125 -

126 -

127 -@pytest.mark.parametrize("seed", RANDOM_SEEDS)

128 -@pytest.mark.parametrize("device", CUDA_DEVICES)

129 -def test_sampler_all_random_seed(seed: int, device: str):

130 - set_random_seed(seed)

131 - torch.set_default_device(device)

132 - batch_size = random.randint(1, 256)

133 - _, fake_logits, sampler = _prepare_test(batch_size)

134 -

135 - for i in range(batch_size):

136 - fake_logits[i, i] = 1e2

137 -

138 - sampling_params = SamplingParams(

139 - temperature=1.0,

140 - n=random.randint(1, 10),

141 - seed=random.randint(0, 10000),

142 - )

143 - sampler_output = _do_sample(batch_size, fake_logits, sampler,

144 - sampling_params, device)

145 -

146 - for i, sequence_output in enumerate(sampler_output):

147 - for nth_output in sequence_output.samples:

148 - assert nth_output.output_token == i

149 -

150 -

151 -@pytest.mark.parametrize("seed", RANDOM_SEEDS)

152 -@pytest.mark.parametrize("device", CUDA_DEVICES)

153 -def test_sampler_all_random_seed_deterministic(seed: int, device: str):

154 - set_random_seed(seed)

155 - torch.set_default_device(device)

156 - batch_size = random.randint(1, 256)

157 - _, fake_logits, sampler = _prepare_test(batch_size)

158 -

159 - sampling_params = SamplingParams(

160 - temperature=1.0,

161 - n=random.randint(1, 10),

162 - seed=random.randint(0, 10000),

163 - )

164 - first_sampler_output = _do_sample(batch_size, fake_logits, sampler,

165 - sampling_params, device)

166 -

167 - second_sampler_output = _do_sample(batch_size, fake_logits, sampler,

168 - sampling_params, device)

169 -

170 - assert first_sampler_output == second_sampler_output

171 -

172 -

173 -@pytest.mark.parametrize("seed", RANDOM_SEEDS)

174 -@pytest.mark.parametrize("device", CUDA_DEVICES)

175 -def test_sampler_min_tokens_penalty(seed: int, device: str):

176 - seq_id_counter = Counter(start=random.randint(0, 100))

177 - set_random_seed(seed)

178 - torch.set_default_device(device)

179 -

180 - def create_sampling_params(min_tokens,

181 - eos_token_id=0,

182 - *,

183 - stop_token_ids: Optional[list[int]] = None,

184 - prompt_logprobs: Optional[int] = None):

185 - sampling_params = SamplingParams(

186 - min_tokens=min_tokens,

187 - max_tokens=9999, # keep higher than max of min_tokens

188 - stop_token_ids=stop_token_ids,

189 - # requesting prompt_logprobs changes the structure of `logits`

190 - prompt_logprobs=prompt_logprobs,

191 - )

192 - sampling_params.all_stop_token_ids.add(eos_token_id)

193 - return sampling_params

194 -

195 - def create_sequence_data(num_input=3, num_generated=0):

196 - seq_data = SequenceData.from_seqs(

197 - random.choices(range(0, VOCAB_SIZE), k=num_input))

198 - if num_generated > 0:

199 - seq_data.output_token_ids = random.choices(range(0, VOCAB_SIZE),

200 - k=num_generated)

201 - return seq_data

202 -

203 - def generate_test_case():

204 - # generate multiple seq groups but limit total batch size

205 - batch_size = random.randint(1, 128)

206 -

207 - expected_penalization = []

208 - sequence_metadata_list: list[SequenceGroupMetadata] = []

209 - # 20% chance to generate seq group metadata list with all prompts

210 - is_prompt = random.random() < 0.2

211 - while batch_size > 0:

212 - num_seqs = 1 if is_prompt else random.randint(1, batch_size)

213 -

214 - eos_token_id = random.randint(0, VOCAB_SIZE - 1)

215 - min_tokens = random.randint(0, 50)

216 - num_stop_tokens = random.randint(0, 8)

217 - if num_stop_tokens > 0:

218 - stop_token_ids = random.choices(range(0, VOCAB_SIZE - 1),

219 - k=num_stop_tokens)

220 - else:

221 - stop_token_ids = None

222 -

223 - sampling_params = create_sampling_params(

224 - min_tokens=min_tokens,

225 - eos_token_id=eos_token_id,

226 - stop_token_ids=stop_token_ids)

227 -

228 - seq_data: dict[int, SequenceData] = {}

229 - seq_group_penalization: list[bool] = []

230 - for _ in range(num_seqs):

231 - num_input = random.randint(1, 100)

232 - num_generated = 0 if is_prompt else random.randint(1, 100)

233 - seq_data[next(seq_id_counter)] = create_sequence_data(

234 - num_input=num_input, num_generated=num_generated)

235 - seq_group_penalization.append(num_generated < min_tokens)

236 -

237 - expected_penalization.extend(seq_group_penalization)

238 - sequence_metadata_list.append(

239 - SequenceGroupMetadata(

240 - request_id=f"test_{batch_size}",

241 - is_prompt=is_prompt,

242 - seq_data=seq_data,

243 - sampling_params=sampling_params,

244 - block_tables={},

245 - ))

246 - batch_size -= num_seqs

247 -

248 - return {

249 - "expected_penalization": expected_penalization,

250 - "seq_group_metadata_list": sequence_metadata_list,

251 - }

252 -

253 - # define some explicit test cases for edge case behavior

254 - prompt_without_penalization = {

255 - "expected_penalization": [False],

256 - "seq_group_metadata_list": [

257 - SequenceGroupMetadata(

258 - request_id="test_1",

259 - is_prompt=True,

260 - seq_data={

261 - next(seq_id_counter): create_sequence_data(),

262 - },

263 - sampling_params=create_sampling_params(0),

264 - block_tables={},

265 - ),

266 - ]

267 - }

268 -

269 - prompt_with_penalization = {

270 - "expected_penalization": [True],

271 - "seq_group_metadata_list": [

272 - SequenceGroupMetadata(

273 - request_id="test_1",

274 - is_prompt=True,

275 - seq_data={

276 - next(seq_id_counter): create_sequence_data(),

277 - },

278 - sampling_params=create_sampling_params(1),

279 - block_tables={},

280 - ),

281 - ]

282 - }

283 -

284 - prompt_with_penalization_and_prompt_logprobs = {

285 - "expected_penalization": [False, False, True],

286 - "seq_group_metadata_list": [

287 - SequenceGroupMetadata(

288 - request_id="test_1",

289 - is_prompt=True,

290 - seq_data={

291 - next(seq_id_counter): create_sequence_data(num_input=3),

292 - },

293 - sampling_params=create_sampling_params(1, prompt_logprobs=3),

294 - block_tables={},

295 - ),

296 - ]

297 - }

298 -

299 - stop_penalizing_after_min_tokens = {

300 - "expected_penalization": [False],

301 - "seq_group_metadata_list": [

302 - SequenceGroupMetadata(

303 - request_id="test_1",

304 - is_prompt=False,

305 - seq_data={

306 - next(seq_id_counter):

307 - create_sequence_data(num_generated=1),

308 - },

309 - sampling_params=create_sampling_params(1),

310 - block_tables={},

311 - )

312 - ]

313 - }

314 -

315 - stop_token_ids = [42, 99, 42, 0] # intentional duplication

316 - prompt_combination = {

317 - "expected_penalization": [False, True, False],

318 - "seq_group_metadata_list": [

319 - SequenceGroupMetadata(

320 - request_id="test_2",

321 - is_prompt=True,

322 - seq_data={

323 - next(seq_id_counter): create_sequence_data(num_input=2),

324 - },

325 - sampling_params=create_sampling_params(1, prompt_logprobs=3),

326 - block_tables={},

327 - ),

328 - SequenceGroupMetadata(

329 - request_id="test_3",

330 - is_prompt=True,

331 - seq_data={

332 - next(seq_id_counter): create_sequence_data(),

333 - },

334 - sampling_params=create_sampling_params(

335 - 0, stop_token_ids=stop_token_ids),

336 - block_tables={},

337 - )

338 - ]

339 - }

340 -

341 - stop_token_ids = [1, 999, 37, 37] # intentional duplication

342 - decode_combination = {

343 - "expected_penalization": [True, False, False, True, False],

344 - "seq_group_metadata_list": [

345 - SequenceGroupMetadata(

346 - request_id="test_1",

347 - is_prompt=False,

348 - seq_data={

349 - next(seq_id_counter):

350 - create_sequence_data(num_generated=1),

351 - next(seq_id_counter):

352 - create_sequence_data(num_generated=100),

353 - },

354 - sampling_params=create_sampling_params(

355 - 2, stop_token_ids=stop_token_ids),

356 - block_tables={},

357 - ),

358 - SequenceGroupMetadata(

359 - request_id="test_2",

360 - is_prompt=False,

361 - seq_data={

362 - next(seq_id_counter):

363 - create_sequence_data(num_generated=20),

364 - next(seq_id_counter):

365 - create_sequence_data(num_generated=1),

366 - next(seq_id_counter):

367 - create_sequence_data(num_generated=10),

368 - },

369 - sampling_params=create_sampling_params(

370 - 10, prompt_logprobs=5, stop_token_ids=stop_token_ids),

371 - block_tables={},

372 - ),

373 - ]

374 - }

375 -

376 - if seed == 0:

377 - test_cases = [

378 - prompt_without_penalization,

379 - prompt_with_penalization,

380 - prompt_with_penalization_and_prompt_logprobs,

381 - stop_penalizing_after_min_tokens,

382 - prompt_combination,

383 - decode_combination,

384 - ]

385 - else:

386 - test_cases = [generate_test_case()]

387 -

388 - def run_test_case(*, expected_penalization: list[bool],

389 - seq_group_metadata_list: list[SequenceGroupMetadata]):

390 - assert expected_penalization, \

391 - "Invalid test case, need expected_penalization"

392 - assert seq_group_metadata_list, \

393 - "Invalid test case, need seq_group_metadata_list"

394 -

395 - batch_size = 0

396 - seq_lens: list[int] = []

397 - sampling_params_per_row: list[SamplingParams] = []

398 - for sgm in seq_group_metadata_list:

399 - sampling_params = sgm.sampling_params

400 -

401 - num_rows = len(sgm.seq_data)

402 - if sgm.is_prompt:

403 - # a prompt seq_group has only one sequence

404 - seq_data = next(iter(sgm.seq_data.values()))

405 - prompt_len = seq_data.get_prompt_len()

406 - seq_lens.append(prompt_len)

407 -

408 - assert sgm.sampling_params is not None

409 - if sgm.sampling_params.prompt_logprobs:

410 - # with prompt_logprobs each token in the prompt has a row in

411 - # logits

412 - num_rows = prompt_len

413 -

414 - batch_size += num_rows

415 - sampling_params_per_row.extend(

416 - itertools.repeat(sampling_params, num_rows))

417 -

418 - assert len(

419 - expected_penalization

420 - ) == batch_size, \

421 - ("Invalid test case, expected_penalization does not match computed"

422 - "batch size")

423 -

424 - _, fake_logits, sampler = _prepare_test(batch_size)

425 - sampling_metadata = SamplingMetadata.prepare(

426 - seq_group_metadata_list,

427 - seq_lens=seq_lens if seq_lens else None,

428 - query_lens=seq_lens if seq_lens else [1] * batch_size,

429 - device=device,

430 - pin_memory=is_pin_memory_available())

431 - # the logits tensor is modified in-place by the sampler

432 - _ = sampler(logits=fake_logits, sampling_metadata=sampling_metadata)

433 -

434 - for logits_idx, (should_penalize, sampling_params) in enumerate(

435 - zip(expected_penalization, sampling_params_per_row)):

436 -

437 - tokens_to_check = sampling_params.all_stop_token_ids

438 -

439 - if should_penalize:

440 - for token_id in tokens_to_check:

441 - assert fake_logits[logits_idx, token_id] == -float(

442 - 'inf'

443 - ), f"Expected token {token_id} for logits row {logits_idx}"

444 - " to be penalized"

445 - # no other tokens should be set to -inf

446 - assert torch.count_nonzero(

447 - fake_logits[logits_idx, :] == -float('inf')) == len(

448 - tokens_to_check

449 - ), f"Expected only {len(tokens_to_check)} to be penalized"

450 - else:

451 - # no tokens should be set to -inf

452 - assert torch.count_nonzero(

453 - fake_logits[logits_idx, :] ==

454 - -float('inf')) == 0, "No tokens should have been penalized"

455 -

456 - for test_case in test_cases:

457 - run_test_case(**test_case)

458 -

459 -

460 -@pytest.mark.parametrize("seed", RANDOM_SEEDS)

461 -@pytest.mark.parametrize("device", CUDA_DEVICES)

462 -def test_sampler_mixed(seed: int, device: str):

463 - set_random_seed(seed)

464 - torch.set_default_device(device)

465 - batch_size = random.randint(1, 256)

466 - input_tensor, fake_logits, sampler = _prepare_test(batch_size)

467 -

468 - seq_group_metadata_list: list[SequenceGroupMetadata] = []

469 - expected_tokens: list[Optional[list[int]]] = []

470 - seq_lens: list[int] = []

471 - for i in range(batch_size):

472 - expected: Optional[list[int]] = None

473 - sampling_type = random.randint(0, 2)

474 - if sampling_type == 0:

475 - sampling_params = SamplingParams(temperature=0)

476 - expected = [int(torch.argmax(fake_logits[i], dim=-1).item())]

477 - elif sampling_type in (1, 2):

478 - n = random.randint(1, 10)

479 - sampling_params = SamplingParams(

480 - temperature=random.random() + 0.1,

481 - top_p=min(random.random() + 0.1, 1),

482 - top_k=random.randint(0, 10),

483 - n=n,

484 - presence_penalty=random.randint(0, 1),

485 - )

486 - if sampling_type == 2:

487 - sampling_params.seed = random.randint(0, 10000)

488 - else:

489 - for idx in range(n):

490 - fake_logits[i, i + idx] = 1e2

491 - expected = list(range(i, i + n))

492 -

493 - expected_tokens.append(expected)

494 - seq_group_metadata_list.append(

495 - SequenceGroupMetadata(

496 - request_id=f"test_{i}",

497 - is_prompt=True,

498 - seq_data={0: SequenceData.from_seqs([1, 2, 3])},

499 - sampling_params=sampling_params,

500 - block_tables={0: [1]},

501 - ))

502 - seq_lens.append(seq_group_metadata_list[-1].seq_data[0].get_len())

503 -

504 - generators: dict[str, torch.Generator] = {}

505 -

506 - def test_sampling():

507 - sampling_metadata = SamplingMetadata.prepare(

508 - seq_group_metadata_list,

509 - seq_lens,

510 - query_lens=seq_lens,

511 - device=device,

512 - pin_memory=is_pin_memory_available(),

513 - generators=generators)

514 - sampler_output = sampler(logits=fake_logits,

515 - sampling_metadata=sampling_metadata)

516 -

517 - for i, (sequence_output, metadata) in enumerate(

518 - zip(sampler_output, seq_group_metadata_list)):

519 - assert metadata.sampling_params is not None

520 -

521 - if (metadata.sampling_params.seed is not None

522 - and expected_tokens[i] is None):

523 - # Record seeded random result to compare with results of

524 - # second invocation

525 - expected_tokens[i] = [

526 - nth_output.output_token

527 - for nth_output in sequence_output.samples

528 - ]

529 - continue

530 -

531 - expected_tokens_item = expected_tokens[i]

532 - assert expected_tokens_item is not None

533 -

534 - for n, nth_output in enumerate(sequence_output.samples):

535 - assert metadata.sampling_params is not None

536 -

537 - if (metadata.sampling_params.temperature == 0

538 - or metadata.sampling_params.seed is not None):

539 - # Ensure exact matches for greedy or random with seed

540 - assert nth_output.output_token == expected_tokens_item[n]

541 - else:

542 - # For non-seeded random check that one of the high-logit

543 - # tokens were chosen

544 - assert nth_output.output_token in expected_tokens_item

545 -

546 - # Test batch

547 - test_sampling()

548 -

549 - # Shuffle the batch and resample

550 - target_index = list(range(batch_size))

551 - for list_to_shuffle in (target_index, seq_group_metadata_list,

552 - expected_tokens, seq_lens):

553 - random.Random(seed).shuffle(list_to_shuffle)

554 - target_index = torch.tensor(target_index)

555 - input_tensor.data = input_tensor.index_select(0, target_index)

556 - fake_logits.data = fake_logits.index_select(0, target_index)

557 -

558 - # This time, results of seeded random samples will be compared with

559 - # the corresponding sample in the pre-shuffled batch

560 - test_sampling()

561 -

562 -

563 -@pytest.mark.parametrize("seed", RANDOM_SEEDS)

564 -@pytest.mark.parametrize("device", CUDA_DEVICES)

565 -def test_sampler_top_k_top_p(seed: int, device: str):

566 - set_random_seed(seed)

567 - batch_size = random.randint(1, 256)

568 - top_k = random.randint(100, 500)

569 - top_p = random.random() * 0.1

570 - vocab_size = 32000

571 - input_tensor = torch.rand((batch_size, 1024),

572 - device=device,

573 - dtype=torch.float16)

574 - fake_logits = torch.normal(0,

575 - 5,

576 - size=(batch_size, vocab_size),

577 - device=input_tensor.device,

578 - dtype=input_tensor.dtype)

579 - sampler = MockLogitsSampler(fake_logits)

580 -

581 - generation_model = GenerationMixin()

582 - generation_config = GenerationConfig(top_k=top_k,

583 - top_p=top_p,

584 - do_sample=True)

585 -

586 - @dataclass

587 - class MockConfig:

588 - is_encoder_decoder: bool = False

589 -

590 - generation_model.config = MockConfig() # needed by the following method

591 - generation_model._prepare_special_tokens(generation_config, device=device)

592 - processors = generation_model._get_logits_processor(generation_config,

593 - None,

594 - None,

595 - None, [],

596 - device=device)

597 - assert len(processors) == 2 # top_p and top_k

598 -

599 - seq_group_metadata_list: list[SequenceGroupMetadata] = []

600 - seq_lens: list[int] = []

601 - for i in range(batch_size):

602 - seq_group_metadata_list.append(

603 - SequenceGroupMetadata(

604 - request_id=f"test_{i}",

605 - is_prompt=True,

606 - seq_data={0: SequenceData.from_seqs([1, 2, 3])},

607 - sampling_params=SamplingParams(

608 - temperature=1,

609 - top_k=top_k,

610 - top_p=top_p,

611 - ),

612 - block_tables={0: [1]},

613 - ))

614 - seq_lens.append(seq_group_metadata_list[-1].seq_data[0].get_len())

615 -

616 - sampling_metadata = SamplingMetadata.prepare(

617 - seq_group_metadata_list,

618 - seq_lens,

619 - query_lens=seq_lens,

620 - device=device,

621 - pin_memory=is_pin_memory_available())

622 -

623 - sample_probs = None

624 -

625 - def mock_sample(probs, *args, **kwargs):

626 - nonlocal sample_probs

627 - sample_probs = probs

628 - return ([[prob.topk(1, dim=-1).indices.tolist(), [0]]

629 - for prob in probs], None)

630 -

631 - # top-k and top-p is only calculated when flashinfer kernel is not available

632 - with patch("vllm.model_executor.layers.sampler._sample", mock_sample), \

633 - patch("vllm.model_executor.layers.sampler."

634 - "flashinfer_top_k_top_p_sampling", None):

635 - sampler(logits=fake_logits, sampling_metadata=sampling_metadata)

636 -

637 - assert sample_probs is not None

638 -

639 - hf_probs = processors(torch.zeros_like(fake_logits), fake_logits.clone())

640 - hf_probs = torch.softmax(hf_probs, dim=-1, dtype=torch.float)

641 - torch.testing.assert_close(hf_probs, sample_probs, rtol=0.0, atol=1e-5)

642 - assert torch.equal(hf_probs.eq(0), sample_probs.eq(0))

643 -

644 -

645 -@pytest.mark.parametrize("seed", RANDOM_SEEDS)

646 -@pytest.mark.parametrize("device", CUDA_DEVICES)

647 -def test_flashinfer_fallback(seed: int, device: str):

648 - if not envs.VLLM_USE_FLASHINFER_SAMPLER:

649 - pytest.skip("Flashinfer sampler is disabled")

650 -

651 - pytest.skip("After FlashInfer 0.2.3, sampling will never fail")

652 -

653 - set_random_seed(seed)

654 - torch.set_default_device(device)

655 - batch_size = random.randint(1, 256)

656 - _, fake_logits, sampler = _prepare_test(batch_size)

657 -

658 - def failing_flashinfer_sampling(*_args, **_kwargs):

659 - return None, torch.zeros(batch_size, device=device, dtype=torch.int32)

660 -

661 - sampling_params = SamplingParams(

662 - temperature=1.0,

663 - n=random.randint(1, 10),

664 - seed=random.randint(0, 10000),

665 - )

666 - sampler_output = _do_sample(batch_size, fake_logits, sampler,

667 - sampling_params, device)

668 -

669 - with patch(

670 - "vllm.model_executor.layers.sampler."

671 - "flashinfer_top_k_top_p_sampling", failing_flashinfer_sampling):

672 - fallback_sampler_output = _do_sample(batch_size, fake_logits, sampler,

673 - sampling_params, device)

674 -

675 - assert sampler_output == fallback_sampler_output

676 -

677 -

678 -@pytest.mark.parametrize("device", CUDA_DEVICES)

679 -def test_sampler_repetition_penalty_mixed(device: str):

680 -

681 - vocab_size = 8

682 -

683 - def test_sampling_params(sampling_params: list[SamplingParams]):

684 -

685 - seq_group_metadata_list: list[SequenceGroupMetadata] = []

686 - seq_lens: list[int] = []

687 - for i in range(2):

688 - seq_group_metadata_list.append(

689 - SequenceGroupMetadata(

690 - request_id=f"test_{i}",

691 - is_prompt=True,

692 - seq_data={0: SequenceData.from_seqs([1, 2, 3])},

693 - sampling_params=sampling_params[i],

694 - block_tables={0: [1]},

695 - ))

696 - seq_lens.append(seq_group_metadata_list[-1].seq_data[0].get_len())

697 -

698 - sampling_metadata = SamplingMetadata.prepare(

699 - seq_group_metadata_list,

700 - seq_lens,

701 - query_lens=seq_lens,

702 - device=device,

703 - pin_memory=is_pin_memory_available())

704 -

705 - fake_logits = torch.full((2, vocab_size),

706 - 1e-2,

707 - device=device,

708 - dtype=torch.float16)

709 -

710 - fake_logits[:, 5] = 1.1e-2

711 - fake_logits[:, 1] = 1.2e-2

712 -

713 - sampler = MockLogitsSampler(fake_logits)

714 -

715 - sampler_output = sampler(logits=fake_logits,

716 - sampling_metadata=sampling_metadata)

717 -

718 - generated_tokens = []

719 - for output in sampler_output:

720 - generated_tokens.append(output.samples[0].output_token)

721 -

722 - return generated_tokens

723 -

724 - # one configuration is greedy with repetition_penalty

725 - sampling_params_rep = SamplingParams(

726 - temperature=0.0,

727 - repetition_penalty=2.0,

728 - )

729 -

730 - # other configuration is sampling w/o repetition_penalty

731 - sampling_params_sample = SamplingParams(

732 - temperature=1.0,

733 - top_k=1,

734 - seed=42,

735 - )

736 -

737 - tokens1 = test_sampling_params(

738 - [sampling_params_rep, sampling_params_sample])

739 -

740 - tokens2 = test_sampling_params(

741 - [sampling_params_sample, sampling_params_rep])

742 -

743 - assert tokens1[0] == tokens2[1]

744 - assert tokens1[1] == tokens2[0]

745 -

746 -

747 -@pytest.mark.parametrize("device", CUDA_DEVICES)

748 -def test_sampler_include_gpu_probs_tensor(device: str):

749 - set_random_seed(42)

750 - torch.set_default_device(device)

751 - batch_size = random.randint(1, 256)

752 - _, fake_logits, sampler = _prepare_test(batch_size)

753 - sampler.include_gpu_probs_tensor = True

754 - sampler.should_modify_greedy_probs_inplace = False

755 -

756 - sampling_params = SamplingParams(temperature=0)

757 -

758 - mock_inplace = Mock()

759 - with patch(

760 - "vllm.model_executor.layers.sampler._modify_greedy_probs_inplace",

761 - mock_inplace):

762 -

763 - sampler_output = _do_sample(batch_size, fake_logits, sampler,

764 - sampling_params, device)

765 - mock_inplace.assert_not_called()

766 -

767 - assert sampler_output.sampled_token_probs is not None

768 - assert sampler_output.logprobs is not None

769 - assert sampler_output.sampled_token_ids is not None

tests/samplers/test_seeded_generate.py deleted

-86

@@ -1,86 +0,0 @@

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

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

3 -"""Verify that seeded random sampling is deterministic.

4 -

5 -Run `pytest tests/samplers/test_seeded_generate.py`.

6 -"""

7 -import copy

8 -import random

9 -from itertools import combinations

10 -

11 -import pytest

12 -

13 -from vllm import SamplingParams

14 -from vllm.model_executor.utils import set_random_seed

15 -

16 -MODEL = "facebook/opt-125m"

17 -RANDOM_SEEDS = list(range(5))

18 -

19 -

20 -@pytest.fixture

21 -def vllm_model(vllm_runner, monkeypatch):

22 - # This file relies on V0 internals.

23 - monkeypatch.setenv("VLLM_USE_V1", "0")

24 - with vllm_runner(MODEL, dtype="half") as vllm_model:

25 - yield vllm_model

26 -

27 -

28 -@pytest.mark.parametrize("seed", RANDOM_SEEDS)

29 -def test_random_sample_with_seed(

30 - vllm_model,

31 - example_prompts,

32 - seed: int,

33 -) -> None:

34 - set_random_seed(seed)

35 -

36 - sampling_params = SamplingParams(

37 - # Parameters to ensure sufficient randomness

38 - temperature=3.0,

39 - top_p=min(random.random() + 0.3, 1),

40 - top_k=random.randint(5, 20),

41 - n=random.randint(1, 10),

42 - presence_penalty=random.randint(0, 1),

43 - max_tokens=8,

44 - ignore_eos=True,

45 - )

46 -

47 - sampling_params_seed_1 = copy.deepcopy(sampling_params)

48 - sampling_params_seed_1.seed = 100

49 - sampling_params_seed_2 = copy.deepcopy(sampling_params)

50 - sampling_params_seed_2.seed = 200

51 -

52 - llm = vllm_model.llm

53 -

54 - for prompt in example_prompts:

55 - for params in (

56 - sampling_params,

57 - sampling_params_seed_1,

58 - sampling_params_seed_2,

59 - sampling_params,

60 - sampling_params_seed_1,

61 - sampling_params_seed_2,

62 - ):

63 - llm._add_request(prompt, params=params)

64 -

65 - results = llm._run_engine(use_tqdm=False)

66 - all_outputs = [[out.token_ids for out in output.outputs]

67 - for output in results]

68 -

69 - for i in range(0, len(example_prompts), 6):

70 - outputs = all_outputs[i:i + 6]

71 -

72 - # verify all non-seeded requests differ

73 - for output_a, output_b in combinations(

74 - (outputs[0], outputs[1], outputs[2], outputs[3]),

75 - 2,

76 - ):

77 - assert output_a != output_b

78 -

79 - # verify requests with the same seed match

80 - assert outputs[1] == outputs[4]

81 - assert outputs[2] == outputs[5]

82 -

83 - # verify generations within the same parallel sampling group differ

84 - for output in outputs:

85 - for sub_output_a, sub_output_b in combinations(output, 2):

86 - assert sub_output_a != sub_output_b