[Bugfix][Spec Decode] Fix probabilistic sampling for parallel drafting (#45956)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>

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,

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
         )