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

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 3, 17:45 UTC

[Bugfix][Spec Decode] Wire draft_probs into probabilistic draft_model rejection (#40269)

Siddharth Bedekar committed May 13, 2026 at 18:04 UTC f51f6844f99aa38547f1fcae6516da31997bda50

7 files changed +264 -13

tests/test_config.py

+18

@@ -1286,6 +1286,24 @@ def test_eagle_draft_model_config():

1286 assert draft_model_config.architecture == "EagleLlamaForCausalLM"

1287

1288

1289 +def test_draft_sample_method_probabilistic_is_accepted():

1290 + speculative_config = SpeculativeConfig(

1291 + method="ngram",

1292 + num_speculative_tokens=1,

1293 + draft_sample_method="probabilistic",

1294 + )

1295 + assert speculative_config.draft_sample_method == "probabilistic"

1296 +

1297 +

1298 +def test_draft_sample_method_gumbel_is_rejected():

1299 + with pytest.raises(ValidationError):

1300 + SpeculativeConfig(

1301 + method="ngram",

1302 + num_speculative_tokens=1,

1303 + draft_sample_method="gumbel",

1304 + )

1305 +

1306 +

1307 def test_ir_op_priority_default():

1308 """Test that IR op priority defaults are set correctly."""

1309 from vllm.config.kernel import IrOpPriorityConfig

tests/v1/spec_decode/test_eagle.py

+101

@@ -50,6 +50,8 @@ def _create_proposer(

50 num_speculative_tokens: int,

51 attention_backend: str | None = None,

52 parallel_drafting: bool = False,

53 + rejection_sample_method: str = "standard",

54 + draft_sample_method: str = "greedy",

55 ) -> EagleProposer:

56 # Method-dependent setup

57 if method == "eagle":

@@ -81,6 +83,8 @@ def _create_proposer(

83 method=method,

84 num_speculative_tokens=num_speculative_tokens,

85 parallel_drafting=parallel_drafting,

86 + rejection_sample_method=rejection_sample_method,

87 + draft_sample_method=draft_sample_method,

88 )

89 if parallel_drafting:

90 # Overwrite pard_token to avoid crash during init

@@ -997,6 +1001,103 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):

1001 assert torch.equal(result, expected_tokens)

1002

1003

1004 +def test_propose_stores_probabilistic_draft_probs(monkeypatch):

1005 + device = torch.device(DEVICE_TYPE)

1006 + batch_size = 2

1007 + seq_lens = [5, 3]

1008 + total_tokens = sum(seq_lens)

1009 + num_speculative_tokens = 3

1010 + vocab_size = 8

1011 +

1012 + proposer = _create_proposer(

1013 + "draft_model",

1014 + num_speculative_tokens,

1015 + rejection_sample_method="standard",

1016 + draft_sample_method="probabilistic",

1017 + )

1018 + hidden_size = proposer.hidden_size

1019 + expanded_total_tokens = total_tokens + batch_size

1020 +

1021 + model_mock = mock.MagicMock()

1022 + forward_returns = []

1023 + logits_returns = []

1024 + for step in range(num_speculative_tokens):

1025 + token_count = expanded_total_tokens if step == 0 else batch_size

1026 + forward_returns.append(torch.zeros(token_count, hidden_size, device=device))

1027 + logits = torch.full((batch_size, vocab_size), -10.0, device=device)

1028 + logits[0, step + 1] = 5.0

1029 + logits[1, step + 3] = 4.0

1030 + logits_returns.append(logits)

1031 +

1032 + model_mock.side_effect = forward_returns

1033 + model_mock.compute_logits.side_effect = logits_returns

1034 + proposer.model = model_mock

1035 + proposer._draft_attn_layer_names = {"layer.0"}

1036 +

1037 + def fake_compute_probs(logits, sampling_metadata):

1038 + probs = torch.softmax(logits, dim=-1)

1039 + return probs.argmax(dim=-1), probs

1040 +

1041 + monkeypatch.setattr(

1042 + "vllm.v1.spec_decode.llm_base_proposer.compute_probs_and_sample_next_token",

1043 + fake_compute_probs,

1044 + )

1045 +

1046 + batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=seq_lens)

1047 + common_attn_metadata = create_common_attn_metadata(

1048 + batch_spec,

1049 + block_size=BLOCK_SIZE,

1050 + device=device,

1051 + )

1052 +

1053 + attn_metadata_builder_cls, _ = try_get_attention_backend(

1054 + AttentionBackendEnum.FLASH_ATTN

1055 + )

1056 + attn_metadata_builder = attn_metadata_builder_cls(

1057 + kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),

1058 + layer_names=proposer._draft_attn_layer_names,

1059 + vllm_config=proposer.vllm_config,

1060 + device=device,

1061 + )

1062 + proposer.runner = mock.MagicMock()

1063 + mock_attn_group = mock.MagicMock()

1064 + mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder

1065 + mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)

1066 + mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec

1067 + proposer.draft_attn_groups = [mock_attn_group]

1068 +

1069 + sampling_metadata = mock.MagicMock()

1070 + sampling_metadata.all_greedy = False

1071 +

1072 + result = proposer.propose(

1073 + target_token_ids=torch.randint(0, vocab_size, (total_tokens,), device=device),

1074 + target_positions=torch.cat(

1075 + [

1076 + torch.arange(seq_lens[0], device=device),

1077 + torch.arange(seq_lens[1], device=device),

1078 + ]

1079 + ),

1080 + target_hidden_states=torch.randn(total_tokens, hidden_size, device=device),

1081 + next_token_ids=torch.randint(

1082 + 0, vocab_size, (batch_size,), dtype=torch.int32, device=device

1083 + ),

1084 + token_indices_to_sample=None,

1085 + common_attn_metadata=common_attn_metadata,

1086 + sampling_metadata=sampling_metadata,

1087 + )

1088 +

1089 + assert result.shape == (batch_size, num_speculative_tokens)

1090 +

1091 + draft_probs = proposer.take_last_draft_probs()

1092 + assert draft_probs is not None

1093 + assert draft_probs.shape == (batch_size, num_speculative_tokens, vocab_size)

1094 + for step, expected_logits in enumerate(logits_returns):

1095 + assert torch.allclose(

1096 + draft_probs[:, step, :],

1097 + torch.softmax(expected_logits, dim=-1),

1098 + )

1099 +

1100 +

1101 def test_set_inputs_first_pass_dflash():

1102 """

1103 Test for DFlash set_inputs_first_pass.

tests/v1/worker/test_gpu_model_runner.py

+35

@@ -2,6 +2,7 @@

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

4 from types import SimpleNamespace

5 +from unittest.mock import Mock

7 import numpy as np

8 import pytest

@@ -40,6 +41,7 @@ from vllm.v1.kv_cache_interface import (

41 KVCacheTensor,

42 )

43 from vllm.v1.sample.metadata import SamplingMetadata

44 +from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

45 from vllm.v1.worker.gpu_input_batch import InputBatch

46 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

47 from vllm.v1.worker.utils import AttentionGroup, select_common_block_size

@@ -747,6 +749,39 @@ def test_reload_weights_before_load_model(model_runner):

749 model_runner.reload_weights()

750

751

752 +def test_sample_passes_reordered_draft_probs_to_rejection_sampler():

753 + runner = object.__new__(GPUModelRunner)

754 + runner.use_async_scheduling = False

755 + runner.input_batch = SimpleNamespace(

756 + sampling_metadata=Mock(spec=SamplingMetadata),

757 + update_async_output_token_ids=Mock(),

758 + req_ids=["req_a", "req_b", "req_c"],

759 + )

760 + runner.rejection_sampler = Mock(return_value="sampler_output")

761 + runner.sampler = Mock()

762 + runner._draft_prob_req_ids = ["req_c", "req_a", "req_b"]

763 + runner._draft_probs = torch.arange(3 * 3 * 4, dtype=torch.float32).reshape(3, 3, 4)

764 +

765 + spec_decode_metadata = SpecDecodeMetadata.make_dummy(

766 + [[1, 2], [], [3]],

767 + device=torch.device("cpu"),

768 + )

769 + logits = torch.randn(6, 4)

770 +

771 + output = GPUModelRunner._sample(runner, logits, spec_decode_metadata)

772 +

773 + assert output == "sampler_output"

774 + passed_draft_probs = runner.rejection_sampler.call_args.args[1]

775 + expected_draft_probs = torch.cat(

776 + [

777 + runner._draft_probs[1, :2],

778 + runner._draft_probs[0, :1],

779 + ],

780 + dim=0,

781 + )

782 + assert torch.equal(passed_draft_probs, expected_draft_probs)

783 +

784 +

785 def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config):

786 torch.set_default_dtype(torch.float16)

787 layer_0 = "model.layers.0.self_attn.attn"

vllm/config/speculative.py

+5 -5

@@ -67,7 +67,7 @@ SpeculativeMethod = Literal[

67 NgramGPUTypes,

68 ]

69 RejectionSampleMethod = Literal["standard", "synthetic"]

70 -DraftSampleMethod = Literal["greedy", "gumbel"]

70 +DraftSampleMethod = Literal["greedy", "probabilistic"]

73 @config

@@ -256,10 +256,10 @@ class SpeculativeConfig:

256 draft_sample_method: DraftSampleMethod = "greedy"

257 """How the draft model samples tokens. 'greedy' always picks the argmax

258 token, and the draft probabilities are treated as one-hot during rejection

259 - sampling. 'gumbel' adds Gumbel noise for stochastic sampling, and the full

260 - draft logits are used for the probability ratio test during rejection

261 - sampling. This comes at the cost of additional GPU memory usage. This

262 - parameter currently only applies to Model Runner V2."""

259 + sampling. 'probabilistic' samples stochastically from the draft

260 + distribution and uses the full draft logits for the probability ratio test

261 + during rejection sampling. This comes at the cost of additional GPU memory

262 + usage."""

263

264 def compute_hash(self) -> str:

265 """

vllm/v1/spec_decode/llm_base_proposer.py

+49 -3

@@ -225,6 +225,11 @@ class SpecDecodeBaseProposer:

225 device=device,

226 with_numpy=True,

227 )

228 + self._enable_probabilistic_draft_probs = (

229 + self.speculative_config.rejection_sample_method == "standard"

230 + and self.speculative_config.draft_sample_method == "probabilistic"

231 + )

232 + self._last_draft_probs: torch.Tensor | None = None

233

234 self._slot_mapping_buffer = torch.zeros(

235 self.max_positions,

@@ -389,6 +394,30 @@ class SpecDecodeBaseProposer:

394 return self.model.get_top_tokens(hidden_states)

395 return self.model.compute_logits(hidden_states).argmax(dim=-1)

396

397 + def _sample_from_logits(

398 + self,

399 + logits: torch.Tensor,

400 + sampling_metadata: SamplingMetadata,

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

402 + if not self._enable_probabilistic_draft_probs:

403 + return logits.argmax(dim=-1), None

404 + if sampling_metadata.all_greedy:

405 + return logits.argmax(dim=-1), None

406 + return compute_probs_and_sample_next_token(logits, sampling_metadata)

407 +

408 + def _sample_draft_tokens(

409 + self,

410 + hidden_states: torch.Tensor,

411 + sampling_metadata: SamplingMetadata,

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

413 + if not self._enable_probabilistic_draft_probs or sampling_metadata.all_greedy:

414 + return self._greedy_sample(hidden_states), None

415 + logits = self.model.compute_logits(hidden_states)

416 + return self._sample_from_logits(logits, sampling_metadata)

417 +

418 + def take_last_draft_probs(self) -> torch.Tensor | None:

419 + return self._last_draft_probs

420 +

421 def propose(

422 self,

423 # [num_tokens]

@@ -408,6 +437,7 @@ class SpecDecodeBaseProposer:

437 | list[dict[str, torch.Tensor]]

438 | None = None,

439 ) -> torch.Tensor:

440 + self._last_draft_probs = None

441 batch_size = common_attn_metadata.batch_size()

442

443 if self.method in ("eagle3", "dflash"):

@@ -469,7 +499,13 @@ class SpecDecodeBaseProposer:

499

500 # Early exit if there is only one draft token to be generated.

501 if self.num_speculative_tokens == 1 or self.parallel_drafting:

472 - draft_token_ids = self._greedy_sample(sample_hidden_states)

502 + draft_token_ids, draft_probs = self._sample_draft_tokens(

503 + sample_hidden_states, sampling_metadata

504 + )

505 + if draft_probs is not None:

506 + self._last_draft_probs = draft_probs.view(

507 + -1, self.num_speculative_tokens, draft_probs.shape[-1]

508 + ).contiguous()

509 return draft_token_ids.view(-1, self.num_speculative_tokens)

510

511 if self.uses_mrope:

@@ -484,7 +520,10 @@ class SpecDecodeBaseProposer:

520 # (which read via _get_positions) use the correct values.

521 self.positions[:batch_size] = positions

522

487 - draft_token_ids = self._greedy_sample(sample_hidden_states)

523 + draft_token_ids, draft_probs = self._sample_draft_tokens(

524 + sample_hidden_states, sampling_metadata

525 + )

526 + draft_probs_list = None if draft_probs is None else [draft_probs]

527

528 if self.allowed_attn_types is not None:

529 for group_md in per_group_attn_metadata:

@@ -584,11 +623,18 @@ class SpecDecodeBaseProposer:

623 last_hidden_states, hidden_states = ret_hidden_states

624

625 hidden_states = hidden_states[:batch_size]

587 - draft_token_ids = self._greedy_sample(last_hidden_states[:batch_size])

626 + draft_token_ids, draft_probs = self._sample_draft_tokens(

627 + last_hidden_states[:batch_size], sampling_metadata

628 + )

629 + if draft_probs is not None:

630 + assert draft_probs_list is not None

631 + draft_probs_list.append(draft_probs)

632 draft_token_ids_list.append(draft_token_ids)

633

634 # [batch_size, num_speculative_tokens]

635 draft_token_ids = torch.stack(draft_token_ids_list, dim=1)

636 + if draft_probs_list is not None:

637 + self._last_draft_probs = torch.stack(draft_probs_list, dim=1).contiguous()

638 return draft_token_ids

639

640 def _update_positions_dependent_metadata(

vllm/v1/worker/gpu/spec_decode/eagle/speculator.py

+1 -1

@@ -106,7 +106,7 @@ class EagleSpeculator:

106 )

107

108 self.draft_logits: torch.Tensor | None = None

109 - if self.speculative_config.draft_sample_method == "gumbel":

109 + if self.speculative_config.draft_sample_method == "probabilistic":

110 self.draft_logits = torch.zeros(

111 self.max_num_reqs,

112 self.num_speculative_steps,

vllm/v1/worker/gpu_model_runner.py

+55 -4

@@ -806,6 +806,8 @@ class GPUModelRunner(

806

807 # Cached outputs.

808 self._draft_token_ids: list[list[int]] | torch.Tensor | None = None

809 + self._draft_probs: torch.Tensor | None = None

810 + self._draft_prob_req_ids: list[str] | None = None

811 # N-gram GPU path: async D2H buffer/event for per-request valid draft counts.

812 self._num_valid_draft_tokens: torch.Tensor | None = None

813 self._num_valid_draft_tokens_cpu: torch.Tensor | None = None

@@ -3433,9 +3435,10 @@ class GPUModelRunner(

3435 draft_token_ids_cpu, _ = self._get_draft_token_ids_cpu()

3436 self.input_batch.update_async_spec_token_ids(draft_token_ids_cpu)

3437

3438 + draft_probs = self._get_spec_decode_draft_probs(spec_decode_metadata)

3439 sampler_output = self.rejection_sampler(

3440 spec_decode_metadata,

3438 - None, # draft_probs

3441 + draft_probs,

3442 logits,

3443 sampling_metadata,

3444 )

@@ -4280,6 +4283,8 @@ class GPUModelRunner(

4283 )

4284

4285 self._draft_token_ids = None

4286 + self._draft_probs = None

4287 + self._draft_prob_req_ids = None

4288 self._draft_token_req_ids = None

4289 self.valid_sampled_token_count_gpu = None

4290 self.input_batch.prev_sampled_token_ids = None

@@ -4374,6 +4379,8 @@ class GPUModelRunner(

4379 self._draft_token_ids = torch.zeros(

4380 1, device=self.device, dtype=torch.int32

4381 ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)

4382 + self._draft_probs = None

4383 + self._draft_prob_req_ids = None

4384 self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)

4385

4386 with record_function_or_nullcontext("gpu_model_runner: bookkeep"):

@@ -4599,6 +4606,35 @@ class GPUModelRunner(

4606 sampled_count_event.synchronize()

4607 return counts_cpu[: prev_sampled_token_ids.shape[0]].tolist()

4608

4609 + def _get_spec_decode_draft_probs(

4610 + self, spec_decode_metadata: SpecDecodeMetadata

4611 + ) -> torch.Tensor | None:

4612 + if self._draft_probs is None or self._draft_prob_req_ids is None:

4613 + return None

4614 +

4615 + row_by_req_id = {

4616 + req_id: idx for idx, req_id in enumerate(self._draft_prob_req_ids)

4617 + }

4618 + draft_probs_rows: list[torch.Tensor] = []

4619 + for req_id, num_draft in zip(

4620 + self.input_batch.req_ids, spec_decode_metadata.num_draft_tokens

4621 + ):

4622 + if num_draft == 0:

4623 + continue

4624 + row_idx = row_by_req_id.get(req_id)

4625 + if row_idx is None:

4626 + logger.warning(

4627 + "Missing cached draft probabilities for request %s; "

4628 + "falling back to legacy speculative rejection behavior.",

4629 + req_id,

4630 + )

4631 + return None

4632 + draft_probs_rows.append(self._draft_probs[row_idx, :num_draft])

4633 +

4634 + if not draft_probs_rows:

4635 + return None

4636 + return torch.cat(draft_probs_rows, dim=0).contiguous()

4637 +

4638 def propose_draft_token_ids(

4639 self,

4640 scheduler_output: "SchedulerOutput",

@@ -4614,6 +4650,8 @@ class GPUModelRunner(

4650 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

4651 spec_config = self.speculative_config

4652 assert spec_config is not None

4653 + self._draft_probs = None

4654 + self._draft_prob_req_ids = None

4655 if spec_config.method == "ngram":

4656 from vllm.v1.spec_decode.ngram_proposer import NgramProposer

4657

@@ -4859,6 +4897,11 @@ class GPUModelRunner(

4897 num_rejected_tokens_gpu=num_rejected_tokens_gpu,

4898 slot_mappings=slot_mappings,

4899 )

4900 + if hasattr(self.drafter, "take_last_draft_probs"):

4901 + draft_probs = self.drafter.take_last_draft_probs()

4902 + if draft_probs is not None:

4903 + self._draft_probs = draft_probs

4904 + self._draft_prob_req_ids = self.input_batch.req_ids.copy()

4905

4906 return draft_token_ids

4907

@@ -5806,10 +5849,18 @@ class GPUModelRunner(

5849 )

5850

5851 num_tokens = sum(len(ids) for ids in draft_token_ids)

5809 - # draft_probs = torch.randn(

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

5811 - # dtype=logits.dtype)

5852 draft_probs = None

5853 + if (

5854 + self.speculative_config.rejection_sample_method == "standard"

5855 + and self.speculative_config.draft_sample_method == "probabilistic"

5856 + ):

5857 + draft_probs = torch.rand(

5858 + num_tokens,

5859 + logits.shape[-1],

5860 + device=self.device,

5861 + dtype=torch.float32,

5862 + )

5863 + draft_probs = torch.softmax(draft_probs, dim=-1)

5864 logits = torch.randn(

5865 num_tokens + num_reqs,

5866 logits.shape[-1],