diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index fecb72800e..b62d7e90da 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -1074,6 +1074,7 @@ def test_propose_stores_probabilistic_draft_probs(attn_backend, monkeypatch): sampling_metadata = mock.MagicMock() sampling_metadata.all_greedy = False + sampling_metadata.temperature = torch.ones(batch_size, device=device) result = proposer.propose( num_speculative_tokens=num_speculative_tokens, diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 9f46cbd242..c78d066066 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import dataclasses from importlib.util import find_spec from typing import Any, cast @@ -423,6 +424,21 @@ class SpecDecodeBaseProposer: return logits.argmax(dim=-1), None if sampling_metadata.all_greedy: return logits.argmax(dim=-1), None + + # Parallel drafting (e.g. DFlash) samples num_speculative_tokens rows + # per request in a single pass, so logits has batch_size * K rows while + # the sampling metadata is per-request. The rows are request-major + # (K consecutive slots per request), so repeat_interleave the + # per-request temperature to match before probabilistic sampling. + temperature = sampling_metadata.temperature + if temperature is not None and temperature.shape[0] != logits.shape[0]: + assert logits.shape[0] % temperature.shape[0] == 0 + factor = logits.shape[0] // temperature.shape[0] + sampling_metadata = dataclasses.replace( + sampling_metadata, + temperature=temperature.repeat_interleave(factor, dim=0), + ) + return compute_probs_and_sample_next_token( logits, sampling_metadata, self.use_fp64_gumbel )