[Bugfix][Spec Decode] Fix probabilistic sampling for parallel drafting (#45956)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Jun 24, 2026 at 01:36 UTC
4c5bc41ba61640a9946e457afcfddfb5e99a1fb8
2 files changed
+17
tests/v1/spec_decode/test_eagle.py
+1
@@ -1074,6 +1074,7 @@ def test_propose_stores_probabilistic_draft_probs(attn_backend, monkeypatch):
1074
1075
sampling_metadata = mock.MagicMock()
1076
sampling_metadata.all_greedy = False
1077
+ sampling_metadata.temperature = torch.ones(batch_size, device=device)
1078
1079
result = proposer.propose(
1080
num_speculative_tokens=num_speculative_tokens,
vllm/v1/spec_decode/llm_base_proposer.py
+16
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import dataclasses
4
from importlib.util import find_spec
5
from typing import Any, cast
6
@@ -423,6 +424,21 @@ class SpecDecodeBaseProposer:
424
return logits.argmax(dim=-1), None
425
if sampling_metadata.all_greedy:
426
return logits.argmax(dim=-1), None
427
+
428
+ # Parallel drafting (e.g. DFlash) samples num_speculative_tokens rows
429
+ # per request in a single pass, so logits has batch_size * K rows while
430
+ # the sampling metadata is per-request. The rows are request-major
431
+ # (K consecutive slots per request), so repeat_interleave the
432
+ # per-request temperature to match before probabilistic sampling.
433
+ temperature = sampling_metadata.temperature
434
+ if temperature is not None and temperature.shape[0] != logits.shape[0]:
435
+ assert logits.shape[0] % temperature.shape[0] == 0
436
+ factor = logits.shape[0] // temperature.shape[0]
437
+ sampling_metadata = dataclasses.replace(
438
+ sampling_metadata,
439
+ temperature=temperature.repeat_interleave(factor, dim=0),
440
+ )
441
+
442
return compute_probs_and_sample_next_token(
443
logits, sampling_metadata, self.use_fp64_gumbel
444
)