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

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

[V1][spec decode] return logprobs for spec decoding (#26060)

Signed-off-by: Giancarlo Delfin <gdelfin@meta.com> Co-authored-by: Nick Hill <nhill@redhat.com>

Giancarlo Delfin committed Oct 22, 2025 at 22:59 UTC 6644796bf445ad6fb34b25c6efe297ce8fe31f81

8 files changed +393 -187

tests/v1/sample/test_logprobs.py

+94

@@ -2,12 +2,14 @@

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

4 import itertools

5 +import math

6 from collections.abc import Generator

7 from typing import get_args

9 import pytest

10 import torch

12 +from tests.utils import large_gpu_mark

13 from tests.v1.sample.utils import (

14 BatchLogprobsComposition,

15 BatchLogprobsSpecType,

@@ -17,6 +19,7 @@ from tests.v1.sample.utils import (

19 )

20 from vllm import SamplingParams

21 from vllm.config.model import LogprobsMode

22 +from vllm.distributed import cleanup_dist_env_and_memory

24 from ...conftest import HfRunner, VllmRunner

@@ -508,3 +511,94 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode):

511 if logprobs_mode in ("raw_logits", "processed_logits"):

512 assert positive_values > 0

513 del llm

514 +

515 +

516 +@pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode))

517 +@pytest.mark.parametrize(

518 + "model_setup",

519 + [

520 + pytest.param(

521 + (

522 + "eagle",

523 + "meta-llama/Llama-3.1-8B-Instruct",

524 + "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",

525 + ),

526 + marks=large_gpu_mark(min_gb=32),

527 + ),

528 + ],

529 +)

530 +def test_spec_decode_logprobs(

531 + logprobs_mode: LogprobsMode,

532 + model_setup: tuple[str, str, str],

533 + monkeypatch: pytest.MonkeyPatch,

534 +):

535 + """Spec decode logprobs should match those of the base model.

536 +

537 + Args:

538 + logprobs_mode: logprobs mode.

539 + model_setup: Spec decode method, base model name, and

540 + draft model name.

541 + """

542 + from vllm import LLM

543 +

544 + with monkeypatch.context() as m:

545 + m.setenv("VLLM_USE_V1", "1")

546 + prompt = "Hello world"

547 + sampling_params = SamplingParams(

548 + temperature=0, logprobs=3, max_tokens=10, ignore_eos=False

549 + )

550 + method, model_name, spec_model_name = model_setup

551 + max_model_len = 256

552 +

553 + # Run base LLM.

554 + ref_llm = LLM(

555 + model=model_name,

556 + max_logprobs=5,

557 + max_model_len=max_model_len,

558 + seed=42,

559 + logprobs_mode=logprobs_mode,

560 + gpu_memory_utilization=0.4,

561 + )

562 + ref_results = ref_llm.generate([prompt], sampling_params)

563 + # Collect logprobs outputs from reference LLM.

564 + ref_logprobs = []

565 + for output in ref_results[0].outputs:

566 + for logprobs in output.logprobs:

567 + for token_id in logprobs:

568 + ref_logprobs.append(logprobs[token_id])

569 + del ref_llm

570 + torch.cuda.empty_cache()

571 + cleanup_dist_env_and_memory()

572 +

573 + # Run spec decode LLM.

574 + spec_llm = LLM(

575 + model_name,

576 + speculative_config={

577 + "method": method,

578 + "model": spec_model_name,

579 + "num_speculative_tokens": 3,

580 + "max_model_len": max_model_len,

581 + },

582 + max_logprobs=5,

583 + max_model_len=max_model_len,

584 + seed=42,

585 + logprobs_mode=logprobs_mode,

586 + gpu_memory_utilization=0.4,

587 + )

588 + spec_results = spec_llm.generate([prompt], sampling_params)

589 + # Collect logprobs outputs from spec decode LLM.

590 + spec_logprobs = []

591 + for output in spec_results[0].outputs:

592 + for logprobs in output.logprobs:

593 + for token_id in logprobs:

594 + spec_logprobs.append(logprobs[token_id])

595 + del spec_llm

596 + torch.cuda.empty_cache()

597 + cleanup_dist_env_and_memory()

598 +

599 + # Per-token logprobs are expected to be the same.

600 + assert len(ref_logprobs) == len(spec_logprobs)

601 + for ref_logprob, spec_logprob in zip(ref_logprobs, spec_logprobs):

602 + assert math.isclose(ref_logprob.logprob, spec_logprob.logprob, abs_tol=1e-3)

603 + assert ref_logprob.rank == spec_logprob.rank

604 + assert ref_logprob.decoded_token == spec_logprob.decoded_token

tests/v1/sample/test_rejection_sampler.py

+85 -79

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 from typing import Any

4 +from unittest.mock import Mock

6 import pytest

7 import torch

@@ -11,6 +12,7 @@ from vllm.platforms import current_platform

12 from vllm.v1.sample.logits_processor import LogitsProcessors

13 from vllm.v1.sample.metadata import SamplingMetadata

14 from vllm.v1.sample.rejection_sampler import PLACEHOLDER_TOKEN_ID, RejectionSampler

15 +from vllm.v1.sample.sampler import Sampler, SamplerOutput

16 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

18 DEVICE = current_platform.device_type

@@ -18,7 +20,28 @@ DEVICE = current_platform.device_type

21 @pytest.fixture

22 def rejection_sampler():

21 - return RejectionSampler()

23 + mock_sampler = Mock(spec=Sampler)

24 + mock_sampler.logprobs_mode = "raw_logprobs"

25 + return RejectionSampler(mock_sampler)

26 +

27 +

28 +def mock_sampler_output(

29 + rejection_sampler: RejectionSampler, bonus_token_ids: torch.Tensor

30 +):

31 + rejection_sampler.sampler.return_value = SamplerOutput(

32 + sampled_token_ids=bonus_token_ids, logprobs_tensors=None

33 + )

34 +

35 +

36 +def create_spec_decode_metadata(

37 + spec_tokens: list[list[int]], logits: torch.Tensor

38 +) -> SpecDecodeMetadata:

39 + metadata = SpecDecodeMetadata.make_dummy(spec_tokens, device=logits.device)

40 + metadata.target_logits_indices = torch.arange(logits.shape[0])

41 + # Output bonus token ids are mocked, so the bonus logit indices should

42 + # be empty.

43 + metadata.bonus_logits_indices = torch.empty(0, dtype=torch.int32)

44 + return metadata

47 def create_logits_tensor(

@@ -111,19 +134,17 @@ def test_perfect_match(rejection_sampler):

134 metadata = create_sampling_metadata(all_greedy=True)

135 logits = create_logits_tensor(output_tokens)

136 bonus_token_tensor = torch.tensor([output_tokens[0][-1]], device=logits.device)

114 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

115 - spec_tokens, device=logits.device

116 - )

137 + spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)

138

139 + mock_sampler_output(rejection_sampler, bonus_token_tensor)

140 output = rejection_sampler(

141 spec_decode_metadata,

142 draft_probs=None,

121 - target_logits=logits,

122 - bonus_token_ids=bonus_token_tensor,

143 + logits=logits,

144 sampling_metadata=metadata,

145 )

146 expected = torch.tensor([[1, 2, 3, 4]], dtype=torch.int, device=logits.device)

126 - assert torch.equal(output, expected)

147 + assert torch.equal(output.sampled_token_ids, expected)

148

149

150 def test_early_mismatch(rejection_sampler):

@@ -134,15 +155,13 @@ def test_early_mismatch(rejection_sampler):

155 metadata = create_sampling_metadata(all_greedy=True)

156 logits = create_logits_tensor(output_tokens)

157 bonus_token_tensor = torch.tensor([output_tokens[0][-1]], device=logits.device)

137 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

138 - spec_tokens, device=logits.device

139 - )

158 + spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)

159

160 + mock_sampler_output(rejection_sampler, bonus_token_tensor)

161 output = rejection_sampler(

162 spec_decode_metadata,

163 draft_probs=None,

144 - target_logits=logits,

145 - bonus_token_ids=bonus_token_tensor,

164 + logits=logits,

165 sampling_metadata=metadata,

166 )

167 expected = torch.tensor(

@@ -150,7 +169,7 @@ def test_early_mismatch(rejection_sampler):

169 dtype=torch.int,

170 device=logits.device,

171 )

153 - assert torch.equal(output, expected)

172 + assert torch.equal(output.sampled_token_ids, expected)

173

174

175 def test_multiple_sequences(rejection_sampler):

@@ -163,21 +182,19 @@ def test_multiple_sequences(rejection_sampler):

182 bonus_token_tensor = torch.tensor(

183 [output_tokens[0][-1], output_tokens[1][-1]], device=logits.device

184 )

166 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

167 - spec_tokens, device=logits.device

168 - )

185 + spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)

186

187 + mock_sampler_output(rejection_sampler, bonus_token_tensor)

188 output = rejection_sampler(

189 spec_decode_metadata,

190 draft_probs=None,

173 - target_logits=logits,

174 - bonus_token_ids=bonus_token_tensor,

191 + logits=logits,

192 sampling_metadata=metadata,

193 )

194 expected = torch.tensor(

195 [[1, 2, 5], [3, 4, PLACEHOLDER_TOKEN_ID]], dtype=torch.int, device=logits.device

196 )

180 - assert torch.equal(output, expected)

197 + assert torch.equal(output.sampled_token_ids, expected)

198

199

200 def test_single_token_sequence(rejection_sampler):

@@ -188,19 +205,17 @@ def test_single_token_sequence(rejection_sampler):

205 metadata = create_sampling_metadata(all_greedy=True)

206 logits = create_logits_tensor(output_tokens)

207 bonus_token_tensor = torch.tensor([output_tokens[0][-1]], device=logits.device)

191 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

192 - spec_tokens, device=logits.device

193 - )

208 + spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)

209

210 + mock_sampler_output(rejection_sampler, bonus_token_tensor)

211 output = rejection_sampler(

212 spec_decode_metadata,

213 draft_probs=None,

198 - target_logits=logits,

199 - bonus_token_ids=bonus_token_tensor,

214 + logits=logits,

215 sampling_metadata=metadata,

216 )

217 expected = torch.tensor([[1, 2]], dtype=torch.int, device=logits.device)

203 - assert torch.equal(output, expected)

218 + assert torch.equal(output.sampled_token_ids, expected)

219

220

221 def test_empty_sequence(rejection_sampler):

@@ -211,19 +226,17 @@ def test_empty_sequence(rejection_sampler):

226 metadata = create_sampling_metadata(all_greedy=True)

227 logits = create_logits_tensor(output_tokens)

228 bonus_token_tensor = torch.tensor([output_tokens[0][-1]], device=logits.device)

214 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

215 - spec_tokens, device=logits.device

216 - )

229 + spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)

230

231 + mock_sampler_output(rejection_sampler, bonus_token_tensor)

232 output = rejection_sampler(

233 spec_decode_metadata,

234 draft_probs=None,

221 - target_logits=logits,

222 - bonus_token_ids=bonus_token_tensor,

235 + logits=logits,

236 sampling_metadata=metadata,

237 )

238 expected = torch.tensor([[5]], dtype=torch.int, device=logits.device)

226 - assert torch.equal(output, expected)

239 + assert torch.equal(output.sampled_token_ids, expected)

240

241

242 def test_multiple_mismatches(rejection_sampler):

@@ -236,15 +249,13 @@ def test_multiple_mismatches(rejection_sampler):

249 bonus_token_tensor = torch.tensor(

250 [output_tokens[0][-1], output_tokens[1][-1]], device=logits.device

251 )

239 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

240 - spec_tokens, device=logits.device

241 - )

252 + spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)

253

254 + mock_sampler_output(rejection_sampler, bonus_token_tensor)

255 output = rejection_sampler(

256 spec_decode_metadata,

257 draft_probs=None,

246 - target_logits=logits,

247 - bonus_token_ids=bonus_token_tensor,

258 + logits=logits,

259 sampling_metadata=metadata,

260 )

261 expected = torch.tensor(

@@ -255,7 +266,7 @@ def test_multiple_mismatches(rejection_sampler):

266 dtype=torch.int,

267 device=logits.device,

268 )

258 - assert torch.equal(output, expected)

269 + assert torch.equal(output.sampled_token_ids, expected)

270

271

272 @pytest.mark.parametrize(

@@ -277,19 +288,17 @@ def test_parametrized_cases(rejection_sampler, spec_tokens, output_tokens, expec

288 bonus_token_tensor = torch.tensor(

289 [tokens[-1] for tokens in output_tokens], device=logits.device

290 )

280 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

281 - spec_tokens, device=logits.device

282 - )

291 + spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)

292

293 + mock_sampler_output(rejection_sampler, bonus_token_tensor)

294 output = rejection_sampler(

295 spec_decode_metadata,

296 draft_probs=None,

287 - target_logits=logits,

288 - bonus_token_ids=bonus_token_tensor,

297 + logits=logits,

298 sampling_metadata=metadata,

299 )

300 expected_tensor = torch.tensor(expected, dtype=torch.int, device=logits.device)

292 - assert torch.equal(output, expected_tensor)

301 + assert torch.equal(output.sampled_token_ids, expected_tensor)

302

303

304 ########################### Tests for Random Sampling ###################

@@ -331,18 +340,19 @@ def test_deterministic_when_seeded(

340 sampling_metadata = create_sampling_metadata(

341 all_greedy=False, temperature=temperature, generators=seeded_seqs

342 )

334 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

335 - draft_token_ids.tolist(), device=DEVICE

343 + spec_decode_metadata = create_spec_decode_metadata(

344 + draft_token_ids.tolist(), target_logits

345 )

346 +

347 + mock_sampler_output(rejection_sampler, bonus_token_ids)

348 rep_result = rejection_sampler(

349 spec_decode_metadata,

339 - draft_probs=draft_probs,

340 - target_logits=target_logits,

341 - bonus_token_ids=bonus_token_ids,

350 + draft_probs=None,

351 + logits=target_logits,

352 sampling_metadata=sampling_metadata,

353 )

354

345 - results.append(rep_result)

355 + results.append(rep_result.sampled_token_ids)

356

357 for i in range(batch_size):

358 if seeded_mask[i]:

@@ -460,7 +470,9 @@ def estimate_rejection_sampling_pdf(

470 Returns:

471 Estimated probability distribution of the output tokens.

472 """

463 - rejection_sampler = RejectionSampler()

473 + mock_sampler = Mock(spec=Sampler)

474 + mock_sampler.logprobs_mode = "raw_logprobs"

475 + rejection_sampler = RejectionSampler(mock_sampler)

476 num_tokens = num_samples * k

477 # Repeat draft probs num_samples * k times.

478 draft_probs = draft_probs.reshape(1, 1, vocab_size).repeat(num_samples, k, 1)

@@ -483,17 +495,18 @@ def estimate_rejection_sampling_pdf(

495 sampling_metadata = create_sampling_metadata(

496 all_greedy=False, temperature=temperature

497 )

486 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

487 - draft_token_ids.tolist(), device=bonus_token_ids.device

498 + spec_decode_metadata = create_spec_decode_metadata(

499 + draft_token_ids.tolist(), target_logits

500 )

489 - output_token_ids = rejection_sampler(

501 +

502 + mock_sampler_output(rejection_sampler, bonus_token_ids)

503 + sampler_output = rejection_sampler(

504 spec_decode_metadata,

505 draft_probs=draft_probs,

492 - target_logits=target_logits,

493 - bonus_token_ids=bonus_token_ids,

506 + logits=target_logits,

507 sampling_metadata=sampling_metadata,

508 )

496 - output_token_ids = output_token_ids[:, :-1].flatten()

509 + output_token_ids = sampler_output.sampled_token_ids[:, :-1].flatten()

510

511 hist = torch.histogram(

512 output_token_ids.to(dtype=torch.float, device="cpu"),

@@ -532,22 +545,19 @@ def _test_masked_logits(

545 bonus_token_ids = torch.zeros((batch_size, 1), dtype=torch.int64, device=DEVICE)

546

547 # Create spec decode metadata

535 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

536 - draft_token_ids,

537 - device=DEVICE,

538 - )

548 + spec_decode_metadata = create_spec_decode_metadata(draft_token_ids, target_logits)

549

550 # Run rejection sampling

541 - output_token_ids = rejection_sampler(

551 + mock_sampler_output(rejection_sampler, bonus_token_ids)

552 + output = rejection_sampler(

553 spec_decode_metadata,

554 draft_probs=draft_probs,

544 - target_logits=target_logits,

545 - bonus_token_ids=bonus_token_ids,

555 + logits=target_logits,

556 sampling_metadata=sampling_metadata,

557 )

558

559 # Remove bonus tokens and reshape

550 - output_token_ids = output_token_ids[:, :-1].flatten().tolist()

560 + output_token_ids = output.sampled_token_ids[:, :-1].flatten().tolist()

561

562 # Check that all sampled tokens are within the unmasked indices.

563 for i in range(num_tokens):

@@ -665,11 +675,11 @@ def test_frequency_penalties(rejection_sampler):

675 spec_decode_metadata = SpecDecodeMetadata.make_dummy(

676 spec_tokens, device=logits.device

677 )

678 + mock_sampler_output(rejection_sampler, bonus_token_tensor)

679 output = rejection_sampler(

680 spec_decode_metadata,

681 draft_probs=None,

671 - target_logits=logits,

672 - bonus_token_ids=bonus_token_tensor,

682 + logits=logits,

683 sampling_metadata=metadata,

684 )

685 expected = torch.tensor(

@@ -677,7 +687,7 @@ def test_frequency_penalties(rejection_sampler):

687 dtype=torch.int,

688 device=logits.device,

689 )

680 - assert torch.equal(output, expected)

690 + assert torch.equal(output.sampled_token_ids, expected)

691

692

693 def test_bad_words(rejection_sampler):

@@ -707,14 +717,12 @@ def test_bad_words(rejection_sampler):

717 bonus_token_tensor = torch.tensor(

718 [output_tokens[i][-1] for i in range(len(output_tokens))], device=logits.device

719 )

710 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

711 - spec_tokens, device=logits.device

712 - )

720 + spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)

721 + mock_sampler_output(rejection_sampler, bonus_token_tensor)

722 output = rejection_sampler(

723 spec_decode_metadata,

724 draft_probs=None,

716 - target_logits=logits,

717 - bonus_token_ids=bonus_token_tensor,

725 + logits=logits,

726 sampling_metadata=metadata,

727 )

728

@@ -723,7 +731,7 @@ def test_bad_words(rejection_sampler):

731 dtype=torch.int,

732 device=logits.device,

733 )

726 - assert torch.equal(output, expected)

734 + assert torch.equal(output.sampled_token_ids, expected)

735

736

737 def test_allowed_token_ids(rejection_sampler):

@@ -756,14 +764,12 @@ def test_allowed_token_ids(rejection_sampler):

764 bonus_token_tensor = torch.tensor(

765 [output_tokens[i][-1] for i in range(len(output_tokens))], device=logits.device

766 )

759 - spec_decode_metadata = SpecDecodeMetadata.make_dummy(

760 - spec_tokens, device=logits.device

761 - )

767 + spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)

768 + mock_sampler_output(rejection_sampler, bonus_token_tensor)

769 output = rejection_sampler(

770 spec_decode_metadata,

771 draft_probs=None,

765 - target_logits=logits,

766 - bonus_token_ids=bonus_token_tensor,

772 + logits=logits,

773 sampling_metadata=metadata,

774 )

775

@@ -772,4 +778,4 @@ def test_allowed_token_ids(rejection_sampler):

778 dtype=torch.int,

779 device=logits.device,

780 )

775 - assert torch.equal(output, expected)

781 + assert torch.equal(output.sampled_token_ids, expected)

vllm/v1/engine/logprobs.py

+1 -1

@@ -66,7 +66,7 @@ class LogprobsProcessor:

66 assert self.logprobs is not None

67 assert self.cumulative_logprob is not None

69 - token_ids_lst, logprobs_lst, ranks_lst = logprobs_lists

69 + token_ids_lst, logprobs_lst, ranks_lst, _ = logprobs_lists

71 for rank, logprobs, token_ids in zip(ranks_lst, logprobs_lst, token_ids_lst):

72 # Detokenize (non-incrementally).

vllm/v1/outputs.py

+24 -9

@@ -14,34 +14,49 @@ else:

16 class LogprobsLists(NamedTuple):

17 - # [num_reqs, max_num_logprobs + 1]

17 + # [num_reqs x num_generated_tokens, max_num_logprobs + 1]

18 logprob_token_ids: list[list[int]]

19 - # [num_reqs, max_num_logprobs + 1]

19 + # [num_reqs x num_generated_tokens, max_num_logprobs + 1]

20 logprobs: list[list[float]]

21 - # [num_reqs]

21 + # [num_reqs x num_generated_tokens]

22 sampled_token_ranks: list[int]

23 -

24 - def slice(self, start: int, end: int):

23 + # [num_reqs]

24 + # Used for slicing the logprobs in cases like speculative

25 + # decoding where the number of generated tokens may be

26 + # different for each request.

27 + cu_num_generated_tokens: list[int] | None = None

28 +

29 + def slice(self, start_req_idx: int, end_req_idx: int):

30 + if self.cu_num_generated_tokens:

31 + start = self.cu_num_generated_tokens[start_req_idx]

32 + end = self.cu_num_generated_tokens[end_req_idx]

33 + else:

34 + start = start_req_idx

35 + end = end_req_idx

36 return LogprobsLists(

37 self.logprob_token_ids[start:end],

38 self.logprobs[start:end],

39 self.sampled_token_ranks[start:end],

40 + self.cu_num_generated_tokens[start_req_idx:end_req_idx]

41 + if self.cu_num_generated_tokens

42 + else None,

43 )

46 class LogprobsTensors(NamedTuple):

33 - # [num_reqs, max_num_logprobs + 1]

47 + # [num_reqs x num_generated_tokens, max_num_logprobs + 1]

48 logprob_token_ids: torch.Tensor

35 - # [num_reqs, max_num_logprobs + 1]

49 + # [num_reqs x num_generated_tokens, max_num_logprobs + 1]

50 logprobs: torch.Tensor

37 - # [num_reqs]

51 + # [num_reqs x num_generated_tokens]

52 selected_token_ranks: torch.Tensor

40 - def tolists(self):

54 + def tolists(self, cu_num_generated_tokens: list[int] | None = None):

55 return LogprobsLists(

56 self.logprob_token_ids.tolist(),

57 self.logprobs.tolist(),

58 self.selected_token_ranks.tolist(),

59 + cu_num_generated_tokens,

60 )

62 @staticmethod

vllm/v1/sample/rejection_sampler.py

+121 -39

@@ -1,15 +1,19 @@

1 # SPDX-License-Identifier: Apache-2.0

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

4 +from dataclasses import replace

5 +

6 import torch

7 import torch.nn as nn

9 from vllm.logger import init_logger

10 from vllm.triton_utils import tl, triton

11 +from vllm.v1.outputs import LogprobsTensors, SamplerOutput

12 from vllm.v1.sample.metadata import SamplingMetadata

13 from vllm.v1.sample.ops.bad_words import apply_bad_words_with_drafts

14 from vllm.v1.sample.ops.penalties import apply_all_penalties

15 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p

16 +from vllm.v1.sample.sampler import Sampler

17 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

19 logger = init_logger(__name__)

@@ -44,17 +48,22 @@ class RejectionSampler(nn.Module):

48 output tokens = accepted tokens + recovered tokens + bonus tokens

49 """

51 + def __init__(self, sampler: Sampler):

52 + super().__init__()

53 + self.sampler = sampler

54 + logprobs_mode = self.sampler.logprobs_mode

55 + self.is_processed_logprobs_mode = logprobs_mode.startswith("processed")

56 + self.is_logits_logprobs_mode = logprobs_mode.endswith("logits")

57 +

58 def forward(

59 self,

60 metadata: SpecDecodeMetadata,

61 # [num_tokens, vocab_size]

62 draft_probs: torch.Tensor | None,

52 - # [num_tokens, vocab_size]

53 - target_logits: torch.Tensor,

54 - # [batch_size, 1]

55 - bonus_token_ids: torch.Tensor,

63 + # [num_tokens + batch_size, vocab_size]

64 + logits: torch.Tensor,

65 sampling_metadata: SamplingMetadata,

57 - ) -> torch.Tensor:

66 + ) -> SamplerOutput:

67 """

68 Args:

69 metadata:

@@ -63,43 +72,65 @@ class RejectionSampler(nn.Module):

72 Probability distribution for the draft tokens. Shape is

73 [num_tokens, vocab_size]. Can be None if probabilities are

74 not provided, which is the case for ngram spec decode.

66 - target_logits (torch.Tensor):

75 + logits (torch.Tensor):

76 Target model's logits probability distribution.

68 - Shape is [num_tokens, vocab_size]. Here, probabilities from

69 - different requests are flattened into a single tensor because

70 - this is the shape of the output logits.

71 - NOTE: `target_logits` can be updated in place to save memory.

72 - bonus_token_ids (torch.Tensor):

73 - A tensor containing bonus tokens. Shape is [batch_size, 1].

74 - Bonus tokens are added to the end of the sequence if all

75 - proposed tokens are accepted. We generate the bonus tokens

76 - outside of the rejection sampler with the default sampling

77 - strategy. It allows for more flexibility in the sampling

78 - process such as top_p, top_k sampling.

77 + Shape is [num_tokens + batch_size, vocab_size]. Here,

78 + probabilities from different requests are flattened into a

79 + single tensor because this is the shape of the output logits.

80 + NOTE: `logits` can be updated in place to save memory.

81 sampling_metadata (vllm.v1.sample.metadata.SamplingMetadata):

82 Additional metadata needed for sampling, such as temperature,

83 top-k/top-p parameters, or other relevant information.

84 Returns:

83 - output_token_ids (torch.Tensor):

84 - A tensor containing the final output token IDs.

85 + SamplerOutput:

86 + Contains the final output token IDs and their logprobs if

87 + requested.

88 """

89 assert metadata.max_spec_len <= MAX_SPEC_LEN

88 - # Use float32 for the target_logits.

89 - target_logits = target_logits.to(torch.float32)

91 + bonus_logits_indices = metadata.bonus_logits_indices

92 + target_logits_indices = metadata.target_logits_indices

93 +

94 + # When indexing with a tensor (bonus_logits_indices), PyTorch

95 + # creates a new tensor with separate storage from the original

96 + # logits tensor. This means any in-place operations on bonus_logits

97 + # won't affect the original logits tensor.

98 + assert logits is not None

99 + bonus_logits = logits[bonus_logits_indices]

100 + bonus_sampler_output = self.sampler(

101 + logits=bonus_logits,

102 + sampling_metadata=replace(

103 + sampling_metadata,

104 + max_num_logprobs=-1,

105 + ),

106 + predict_bonus_token=True,

107 + # Override the logprobs mode to return logits because they are

108 + # needed later to compute the accepted token logprobs.

109 + logprobs_mode_override="processed_logits"

110 + if self.is_processed_logprobs_mode

111 + else "raw_logits",

112 + )

113 + bonus_token_ids = bonus_sampler_output.sampled_token_ids

114

115 + # Just like `bonus_logits`, `target_logits` is a new tensor with

116 + # separate storage from the original `logits` tensor. Therefore,

117 + # it is safe to update `target_logits` in place.

118 + raw_target_logits = logits[target_logits_indices]

119 + # Use float32 for the target_logits.

120 + raw_target_logits = raw_target_logits.to(torch.float32)

121 target_logits = self.apply_logits_processors(

92 - target_logits, sampling_metadata, metadata

122 + raw_target_logits, sampling_metadata, metadata

123 )

94 -

124 # [num_tokens, vocab_size]

125 # NOTE(woosuk): `target_logits` can be updated in place inside the

97 - # `compute_probs` function.

98 - target_probs = compute_probs(

126 + # `apply_sampling_constraints` function.

127 + target_logits = apply_sampling_constraints(

128 target_logits,

129 metadata.cu_num_draft_tokens,

130 sampling_metadata,

131 )

132 + # Compute probability distribution from target logits.

133 + target_probs = target_logits.softmax(dim=-1, dtype=torch.float32)

134

135 output_token_ids = rejection_sample(

136 metadata.draft_token_ids,

@@ -111,7 +142,63 @@ class RejectionSampler(nn.Module):

142 bonus_token_ids,

143 sampling_metadata,

144 )

114 - return output_token_ids

145 +

146 + logprobs_tensors = None

147 + if sampling_metadata.max_num_logprobs:

148 + logprobs_tensors = self._get_logprobs_tensors(

149 + sampling_metadata.max_num_logprobs,

150 + metadata,

151 + logits,

152 + target_logits if self.is_processed_logprobs_mode else raw_target_logits,

153 + bonus_sampler_output.logprobs_tensors.logprobs,

154 + output_token_ids,

155 + )

156 +

157 + return SamplerOutput(

158 + sampled_token_ids=output_token_ids,

159 + logprobs_tensors=logprobs_tensors,

160 + )

161 +

162 + def _get_logprobs_tensors(

163 + self,

164 + max_num_logprobs: int,

165 + metadata: SpecDecodeMetadata,

166 + logits: torch.Tensor,

167 + target_logits: torch.Tensor,

168 + bonus_logits: torch.Tensor,

169 + sampled_token_ids: torch.Tensor,

170 + ) -> LogprobsTensors:

171 + cu_num_sampled_tokens = torch.zeros_like(metadata.cu_num_sampled_tokens)

172 + cu_num_sampled_tokens[1:] = metadata.cu_num_sampled_tokens[:-1]

173 +

174 + # Collect target and bonus logits.

175 + bonus_logits_indices = metadata.bonus_logits_indices

176 + target_logits_indices = metadata.target_logits_indices

177 + final_logits = torch.zeros_like(logits, dtype=torch.float32)

178 + final_logits[target_logits_indices] = target_logits.to(torch.float32)

179 + final_logits[bonus_logits_indices] = bonus_logits.to(torch.float32)

180 +

181 + # Compute accepted token indices.

182 + accepted_mask = sampled_token_ids != PLACEHOLDER_TOKEN_ID

183 + num_accepted_tokens = accepted_mask.sum(dim=-1)

184 + accepted_logit_indices = accepted_mask.nonzero(as_tuple=True)[1]

185 + accepted_logit_indices += cu_num_sampled_tokens.repeat_interleave(

186 + num_accepted_tokens

187 + )

188 +

189 + # Compute logprobs for accepted tokens.

190 + accepted_logits = final_logits[accepted_logit_indices]

191 + accepted_logprobs = (

192 + accepted_logits

193 + if self.is_logits_logprobs_mode

194 + else self.sampler.compute_logprobs(accepted_logits)

195 + )

196 + accepted_tokens = sampled_token_ids[accepted_mask]

197 + return self.sampler.gather_logprobs(

198 + accepted_logprobs,

199 + max_num_logprobs,

200 + accepted_tokens.to(torch.int64),

201 + )

202

203 @staticmethod

204 def parse_output(

@@ -119,14 +206,12 @@ class RejectionSampler(nn.Module):

206 vocab_size: int,

207 ) -> list[list[int]]:

208 """Parse the output of the rejection sampler.

122 -

209 Args:

210 output_token_ids: The sampled token IDs in shape

211 [batch_size, max_spec_len + 1]. The rejected tokens are

212 replaced with `PLACEHOLDER_TOKEN_ID` by the rejection sampler

213 and will be filtered out in this function.

214 vocab_size: The size of the vocabulary.

129 -

215 Returns:

216 A list of lists of token IDs.

217 """

@@ -328,27 +413,26 @@ def rejection_sample(

413 return output_token_ids

414

415

331 -def compute_probs(

416 +def apply_sampling_constraints(

417 logits: torch.Tensor, # [num_tokens, vocab_size]

418 cu_num_draft_tokens: torch.Tensor, # [batch_size]

419 sampling_metadata: SamplingMetadata,

420 ) -> torch.Tensor:

336 - """Compute probability distribution from logits based on sampling metadata.

421 + """Process logits based on sampling metadata.

422

338 - This function applies temperature scaling to the logits and converts

339 - them to probabilities using softmax. For greedy decoding, it returns

423 + This function applies temperature scaling to the logits,

424 + as well as top-k and top-p. For greedy decoding, it returns

425 the original logits.

426

427 Args:

343 - logits: Input logits tensor to be converted to probabilities.

428 + logits: Input logits tensor to be processed.

429 cu_num_draft_tokens: Cumulative number of draft tokens.

430 sampling_metadata: Metadata containing sampling parameters such as

431 temperature and whether greedy sampling is used.

432

433 Returns:

349 - torch.Tensor: Probability distribution (softmax of scaled logits)

350 - if non-greedy sampling is used, otherwise returns the

351 - original logits.

434 + torch.Tensor: Processed logits if non-greedy sampling is used,

435 + otherwise returns the original logits.

436 """

437 assert logits.ndim == 2

438 assert cu_num_draft_tokens.ndim == 1

@@ -384,9 +468,7 @@ def compute_probs(

468

469 # NOTE(woosuk): `apply_top_k_top_p` uses sorting to calculate the mask,

470 # which is slow for large vocab sizes. This may cause performance issues.

387 - logits = apply_top_k_top_p(logits, top_k, top_p)

388 - output_prob = logits.softmax(dim=-1, dtype=torch.float32)

389 - return output_prob

471 + return apply_top_k_top_p(logits, top_k, top_p)

472

473

474 def expand_batch_to_tokens(

vllm/v1/sample/sampler.py

+20 -11

@@ -69,16 +69,18 @@ class Sampler(nn.Module):

69 logits: torch.Tensor,

70 sampling_metadata: SamplingMetadata,

71 predict_bonus_token: bool = False,

72 + logprobs_mode_override: LogprobsMode | None = None,

73 ) -> SamplerOutput:

74 + logprobs_mode = logprobs_mode_override or self.logprobs_mode

75 # NOTE(woosuk): Use the original logits (before any penalties or

76 # temperature scaling) for the top-k logprobs.

77 # This is different from the V0 sampler, which uses the logits that

78 # is used for sampling (after penalties and temperature scaling).

79 num_logprobs = sampling_metadata.max_num_logprobs

80 if num_logprobs is not None:

79 - if self.logprobs_mode == "raw_logprobs":

81 + if logprobs_mode == "raw_logprobs":

82 raw_logprobs = self.compute_logprobs(logits)

81 - elif self.logprobs_mode == "raw_logits":

83 + elif logprobs_mode == "raw_logits":

84 raw_logprobs = logits.clone()

86 # Use float32 for the logits.

@@ -97,13 +99,18 @@ class Sampler(nn.Module):

99 # return int32 (while PyTorch argmax and topk return int64).

100 sampled = sampled.long()

101

100 - # Gather the logprobs of the topk and sampled token (if requested).

101 - # Get logprobs and rank tensors (if requested)

102 - logprobs_tensors = (

103 - None

104 - if num_logprobs is None

105 - else self.gather_logprobs(raw_logprobs, num_logprobs, token_ids=sampled)

106 - )

102 + if num_logprobs is None:

103 + logprobs_tensors = None

104 + elif num_logprobs == -1:

105 + # Return the full unsorted and unranked logprobs.

106 + logprobs_tensors = LogprobsTensors(

107 + torch.empty(0), raw_logprobs, torch.empty(0)

108 + )

109 + else:

110 + # Gather the logprobs and ranks of the topk and sampled token.

111 + logprobs_tensors = self.gather_logprobs(

112 + raw_logprobs, num_logprobs, token_ids=sampled

113 + )

114

115 # Use int32 to reduce the tensor size.

116 sampled = sampled.to(torch.int32)

@@ -138,6 +145,7 @@ class Sampler(nn.Module):

145 self,

146 logits: torch.Tensor,

147 sampling_metadata: SamplingMetadata,

148 + logprobs_mode_override: LogprobsMode | None = None,

149 ) -> tuple[torch.Tensor, torch.Tensor | None]:

150 """Sample logits based on sampling metadata.

151

@@ -145,6 +153,7 @@ class Sampler(nn.Module):

153 may update the logits tensor in-place.

154 """

155

156 + logprobs_mode = logprobs_mode_override or self.logprobs_mode

157 assert not (sampling_metadata.all_greedy and sampling_metadata.all_random)

158 if sampling_metadata.all_random:

159 greedy_sampled = None

@@ -153,9 +162,9 @@ class Sampler(nn.Module):

162 if sampling_metadata.all_greedy:

163 processed_logprobs = None

164 if sampling_metadata.max_num_logprobs is not None:

156 - if self.logprobs_mode == "processed_logits":

165 + if logprobs_mode == "processed_logits":

166 processed_logprobs = logits

158 - elif self.logprobs_mode == "processed_logprobs":

167 + elif logprobs_mode == "processed_logprobs":

168 processed_logprobs = self.compute_logprobs(logits)

169 return greedy_sampled, processed_logprobs

170

vllm/v1/spec_decode/metadata.py

@@ -14,6 +14,8 @@ class SpecDecodeMetadata:

14 num_draft_tokens: list[int]

15 # [batch_size]

16 cu_num_draft_tokens: torch.Tensor

17 + # [batch_size]

18 + cu_num_sampled_tokens: torch.Tensor

19 # [num_tokens]

20 target_logits_indices: torch.Tensor

21 # [batch_size]

@@ -32,6 +34,7 @@ class SpecDecodeMetadata:

34 ) -> "SpecDecodeMetadata":

35 batch_size = len(draft_token_ids)

36 num_draft_tokens = [len(ids) for ids in draft_token_ids]

37 + num_sampled_tokens = [len(ids) + 1 for ids in draft_token_ids]

38 flattened_draft_token_ids = sum(draft_token_ids, [])

39 num_tokens = len(flattened_draft_token_ids)

@@ -40,6 +43,10 @@ class SpecDecodeMetadata:

43 )

44 cu_num_draft_tokens = np.cumsum(num_draft_tokens, dtype=np.int32)

45 cu_num_draft_tokens_tensor = torch.from_numpy(cu_num_draft_tokens).to(device)

46 + cu_num_sampled_tokens = np.cumsum(num_sampled_tokens, dtype=np.int32)

47 + cu_num_sampled_tokens_tensor = torch.from_numpy(cu_num_sampled_tokens).to(

48 + device

49 + )

51 target_logits_indices = torch.zeros(

52 num_tokens, dtype=torch.int32, device=device

@@ -52,6 +59,7 @@ class SpecDecodeMetadata:

59 draft_token_ids=draft_token_ids_tensor,

60 num_draft_tokens=num_draft_tokens,

61 cu_num_draft_tokens=cu_num_draft_tokens_tensor,

62 + cu_num_sampled_tokens=cu_num_sampled_tokens_tensor,

63 target_logits_indices=target_logits_indices,

64 bonus_logits_indices=bonus_logits_indices,

65 logits_indices=logits_indices,

vllm/v1/worker/gpu_model_runner.py

+40 -48

@@ -327,7 +327,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

327 "Unknown speculative decoding method: "

328 f"{self.speculative_config.method}"

329 )

330 - self.rejection_sampler = RejectionSampler()

330 + self.rejection_sampler = RejectionSampler(self.sampler)

331

332 # Request states.

333 self.requests: dict[str, CachedRequestState] = {}

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

1624 cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to(

1625 self.device, non_blocking=True

1626 )

1627 + cu_num_sampled_tokens = torch.from_numpy(cu_num_sampled_tokens).to(

1628 + self.device, non_blocking=True

1629 + )

1630 logits_indices = torch.from_numpy(logits_indices).to(

1631 self.device, non_blocking=True

1632 )

@@ -1639,15 +1642,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1642 draft_token_ids = self.input_ids.gpu[logits_indices]

1643 draft_token_ids = draft_token_ids[target_logits_indices + 1]

1644

1642 - metadata = SpecDecodeMetadata(

1645 + return SpecDecodeMetadata(

1646 draft_token_ids=draft_token_ids,

1647 num_draft_tokens=num_draft_tokens.tolist(),

1648 cu_num_draft_tokens=cu_num_draft_tokens,

1649 + cu_num_sampled_tokens=cu_num_sampled_tokens,

1650 target_logits_indices=target_logits_indices,

1651 bonus_logits_indices=bonus_logits_indices,

1652 logits_indices=logits_indices,

1653 )

1650 - return metadata

1654

1655 def _prepare_kv_sharing_fast_prefill(

1656 self,

@@ -2221,32 +2224,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2224 sampling_metadata=sampling_metadata,

2225 )

2226

2224 - # When indexing with a tensor (bonus_logits_indices), PyTorch

2225 - # creates a new tensor with separate storage from the original

2226 - # logits tensor. This means any in-place operations on bonus_logits

2227 - # won't affect the original logits tensor.

2228 - assert logits is not None

2229 - bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]

2230 - sampler_output = self.sampler(

2231 - logits=bonus_logits,

2232 - sampling_metadata=sampling_metadata,

2233 - predict_bonus_token=True,

2234 - )

2235 - bonus_token_ids = sampler_output.sampled_token_ids

2236 -

2237 - # Just like `bonus_logits`, `target_logits` is a new tensor with

2238 - # separate storage from the original `logits` tensor. Therefore,

2239 - # it is safe to update `target_logits` in place.

2240 - target_logits = logits[spec_decode_metadata.target_logits_indices]

2241 - output_token_ids = self.rejection_sampler(

2227 + sampler_output = self.rejection_sampler(

2228 spec_decode_metadata,

2229 None, # draft_probs

2244 - target_logits,

2245 - bonus_token_ids,

2230 + logits,

2231 sampling_metadata,

2232 )

2248 - sampler_output.sampled_token_ids = output_token_ids

2249 - self._update_states_after_model_execute(output_token_ids)

2233 + self._update_states_after_model_execute(sampler_output.sampled_token_ids)

2234 return sampler_output

2235

2236 def _bookkeeping_sync(

@@ -2256,6 +2240,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2240 logits: torch.Tensor | None,

2241 hidden_states: torch.Tensor,

2242 num_scheduled_tokens: int,

2243 + spec_decode_metadata: SpecDecodeMetadata | None,

2244 ) -> tuple[

2245 dict[str, int],

2246 LogprobsLists | None,

@@ -2282,19 +2267,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2267 req_ids_output_copy = self.input_batch.req_ids.copy()

2268 req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()

2269

2285 - # NOTE: GPU -> CPU Sync happens here.

2286 - # Move as many CPU operations as possible before this sync point.

2287 - logprobs_tensors = sampler_output.logprobs_tensors

2288 - logprobs_lists = (

2289 - logprobs_tensors.tolists() if logprobs_tensors is not None else None

2290 - )

2291 -

2292 - # Compute prompt logprobs if needed.

2293 - prompt_logprobs_dict = self._get_prompt_logprobs_dict(

2294 - hidden_states[:num_scheduled_tokens],

2295 - scheduler_output.num_scheduled_tokens,

2296 - )

2297 -

2270 num_sampled_tokens = sampler_output.sampled_token_ids.shape[0]

2271 sampled_token_ids = sampler_output.sampled_token_ids

2272 invalid_req_indices = []

@@ -2335,6 +2307,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2307 # the sampled tokens back, because there's no direct communication

2308 # between the first-stage worker and the last-stage worker.

2309 req_ids = self.input_batch.req_ids

2310 + logprobs_tensors = sampler_output.logprobs_tensors

2311 + cu_num_accepted_tokens = (

2312 + [0] if spec_decode_metadata and logprobs_tensors else None

2313 + )

2314 for req_idx in range(num_sampled_tokens):

2315 if self.use_async_scheduling:

2316 sampled_ids = [-1] if req_idx not in invalid_req_indices_set else None

@@ -2360,6 +2336,25 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2336 req_state = self.requests[req_id]

2337 req_state.output_token_ids.extend(sampled_ids)

2338

2339 + if cu_num_accepted_tokens is not None:

2340 + cu_num_accepted_tokens.append(

2341 + cu_num_accepted_tokens[-1] + len(sampled_ids)

2342 + )

2343 +

2344 + # NOTE: GPU -> CPU Sync happens here.

2345 + # Move as many CPU operations as possible before this sync point.

2346 + logprobs_lists = (

2347 + logprobs_tensors.tolists(cu_num_accepted_tokens)

2348 + if logprobs_tensors is not None

2349 + else None

2350 + )

2351 +

2352 + # Compute prompt logprobs if needed.

2353 + prompt_logprobs_dict = self._get_prompt_logprobs_dict(

2354 + hidden_states[:num_scheduled_tokens],

2355 + scheduler_output.num_scheduled_tokens,

2356 + )

2357 +

2358 return (

2359 num_nans_in_logits,

2360 logprobs_lists,

@@ -2644,6 +2639,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2639 logits,

2640 hidden_states,

2641 num_scheduled_tokens,

2642 + spec_decode_metadata,

2643 )

2644

2645 if (

@@ -3560,20 +3556,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3556 # num_tokens, logits.shape[-1], device=self.device,

3557 # dtype=logits.dtype)

3558 draft_probs = None

3563 - target_logits = torch.randn(

3564 - num_tokens, logits.shape[-1], device=self.device, dtype=logits.dtype

3565 - )

3566 - # NOTE(woosuk): Here, we should use int32 because the sampler uses

3567 - # int32 for bonus_token_ids. If the dtype mismatches, re-compilation

3568 - # will occur at runtime.

3569 - bonus_token_ids = torch.zeros(

3570 - num_reqs, device=self.device, dtype=torch.int32

3559 + logits = torch.randn(

3560 + num_tokens + num_reqs,

3561 + logits.shape[-1],

3562 + device=self.device,

3563 + dtype=logits.dtype,

3564 )

3565 self.rejection_sampler(

3566 dummy_spec_decode_metadata,

3567 draft_probs,

3575 - target_logits,

3576 - bonus_token_ids,

3568 + logits,

3569 dummy_metadata,

3570 )

3571 return sampler_output