[V1][spec decode] return logprobs for spec decoding

Author: TheEpicDolphin

tests/v1/sample/test_logprobs.py

@@ -503,3 +503,68 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode,
         if logprobs_mode in ("raw_logits", "processed_logits"):
             assert positive_values > 0
         del llm
+
+@pytest.mark.parametrize("logprobs_mode", ["raw_logprobs"])#get_args(LogprobsMode))
+@pytest.mark.parametrize("model_setup", [("eagle", "meta-llama/Llama-3.1-8B-Instruct", "yuhuili/EAGLE-LLaMA3.1-Instruct-8B")])
+def test_spec_decode_logprobs(
+    logprobs_mode: LogprobsMode,
+    model_setup: tuple[str, str, str],
+    monkeypatch: pytest.MonkeyPatch):
+    """Spec decode logprobs should match those of the base model.
+
+    Args:
+        logprobs_mode: logprobs mode.
+        model_setup: Spec decode method, base model name, and
+        draft model name.
+    """
+    from vllm import LLM
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V1", "1")
+        prompt = "Hello world"
+        sampling_params = SamplingParams(
+            temperature=0,
+            logprobs=1,
+            max_tokens=10,
+            ignore_eos=False)
+        method, model_name, spec_model_name = model_setup
+
+        # Run base LLM.
+        ref_llm = LLM(model=model_name,
+                max_model_len=2048)
+        ref_results = ref_llm.generate([prompt], sampling_params)
+        del ref_llm
+
+        # Run spec decode LLM.
+        spec_llm = LLM(
+            model_name,
+            speculative_config={
+                "method": method,
+                "model": spec_model_name,
+                "num_speculative_tokens": 3,
+                "max_model_len": 2048,
+            },
+            max_logprobs=5,
+            max_model_len=2048,
+            logprobs_mode=logprobs_mode)
+        spec_results = spec_llm.generate([prompt], sampling_params)
+        del spec_llm
+
+        ref_logprobs = []
+        for output in ref_results[0].outputs:
+            for logprobs in output.logprobs:
+                print("[Giancarlo] ", logprobs)
+                for token_id in logprobs:
+                    ref_logprobs.append(logprobs[token_id])
+        print("\n")
+        spec_logprobs = []
+        for output in spec_results[0].outputs:
+            for logprobs in output.logprobs:
+                print("[Giancarlo] ", logprobs)
+                for token_id in logprobs:
+                    spec_logprobs.append(logprobs[token_id])
+        assert len(ref_logprobs) == len(spec_logprobs)
+
+        for ref_logprob, spec_logprob in zip(ref_logprobs, spec_logprobs):
+            assert ref_logprob.logprob == spec_logprob.logprob
+            assert ref_logprob.rank == spec_logprob.rank
+            assert ref_logprob.decoded_token == spec_logprob.decoded_token

tests/v1/sample/test_rejection_sampler.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any, Optional
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -11,14 +12,32 @@
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.rejection_sampler import (PLACEHOLDER_TOKEN_ID,
                                               RejectionSampler)
+from vllm.v1.sample.sampler import Sampler, SamplerOutput
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 
 DEVICE = current_platform.device_type
 
 
 @pytest.fixture
 def rejection_sampler():
-    return RejectionSampler()
+    mock_sampler = Mock(spec=Sampler)
+    return RejectionSampler(mock_sampler)
+
+
+def mock_sampler_output(rejection_sampler: RejectionSampler,
+                             bonus_token_ids: torch.Tensor):
+    rejection_sampler.sampler.return_value = SamplerOutput(
+        sampled_token_ids=bonus_token_ids, logprobs_tensors=None)
+
+
+def create_spec_decode_metadata(spec_tokens: list[list[int]],
+                                logits: torch.Tensor) -> SpecDecodeMetadata:
+    metadata = SpecDecodeMetadata.make_dummy(spec_tokens, device=logits.device)
+    metadata.target_logits_indices = torch.arange(logits.shape[0])
+    # Output bonus token ids are mocked, so the bonus logit indices should
+    # be empty.
+    metadata.bonus_logits_indices = torch.empty(0, dtype=torch.int32)
+    return metadata
 
 
 def create_logits_tensor(output_token_ids: list[list[int]],
@@ -83,20 +102,19 @@ def test_perfect_match(rejection_sampler):
     logits = create_logits_tensor(output_tokens)
     bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
                                       device=logits.device)
-    spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
-                                                         device=logits.device)
+    spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
 
+    mock_sampler_output(rejection_sampler, bonus_token_tensor)
     output = rejection_sampler(
         spec_decode_metadata,
         draft_probs=None,
-        target_logits=logits,
-        bonus_token_ids=bonus_token_tensor,
+        logits=logits,
         sampling_metadata=metadata,
     )
     expected = torch.tensor([[1, 2, 3, 4]],
                             dtype=torch.int,
                             device=logits.device)
-    assert torch.equal(output, expected)
+    assert torch.equal(output.sampled_token_ids, expected)
 
 
 def test_early_mismatch(rejection_sampler):
@@ -108,22 +126,21 @@ def test_early_mismatch(rejection_sampler):
     logits = create_logits_tensor(output_tokens)
     bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
                                       device=logits.device)
-    spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
-                                                         device=logits.device)
+    spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
 
+    mock_sampler_output(rejection_sampler, bonus_token_tensor)
     output = rejection_sampler(
         spec_decode_metadata,
         draft_probs=None,
-        target_logits=logits,
-        bonus_token_ids=bonus_token_tensor,
+        logits=logits,
         sampling_metadata=metadata,
     )
     expected = torch.tensor(
         [[1, 5, PLACEHOLDER_TOKEN_ID, PLACEHOLDER_TOKEN_ID]],
         dtype=torch.int,
         device=logits.device,
     )
-    assert torch.equal(output, expected)
+    assert torch.equal(output.sampled_token_ids, expected)
 
 
 def test_multiple_sequences(rejection_sampler):
@@ -136,20 +153,19 @@ def test_multiple_sequences(rejection_sampler):
     logits = create_logits_tensor(output_tokens)
     bonus_token_tensor = torch.tensor(
         [output_tokens[0][-1], output_tokens[1][-1]], device=logits.device)
-    spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
-                                                         device=logits.device)
+    spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
 
+    mock_sampler_output(rejection_sampler, bonus_token_tensor)
     output = rejection_sampler(
         spec_decode_metadata,
         draft_probs=None,
-        target_logits=logits,
-        bonus_token_ids=bonus_token_tensor,
+        logits=logits,
         sampling_metadata=metadata,
     )
     expected = torch.tensor([[1, 2, 5], [3, 4, PLACEHOLDER_TOKEN_ID]],
                             dtype=torch.int,
                             device=logits.device)
-    assert torch.equal(output, expected)
+    assert torch.equal(output.sampled_token_ids, expected)
 
 
 def test_single_token_sequence(rejection_sampler):
@@ -161,18 +177,17 @@ def test_single_token_sequence(rejection_sampler):
     logits = create_logits_tensor(output_tokens)
     bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
                                       device=logits.device)
-    spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
-                                                         device=logits.device)
+    spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
 
+    mock_sampler_output(rejection_sampler, bonus_token_tensor)
     output = rejection_sampler(
         spec_decode_metadata,
         draft_probs=None,
-        target_logits=logits,
-        bonus_token_ids=bonus_token_tensor,
+        logits=logits,
         sampling_metadata=metadata,
     )
     expected = torch.tensor([[1, 2]], dtype=torch.int, device=logits.device)
-    assert torch.equal(output, expected)
+    assert torch.equal(output.sampled_token_ids, expected)
 
 
 def test_empty_sequence(rejection_sampler):
@@ -184,18 +199,17 @@ def test_empty_sequence(rejection_sampler):
     logits = create_logits_tensor(output_tokens)
     bonus_token_tensor = torch.tensor([output_tokens[0][-1]],
                                       device=logits.device)
-    spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
-                                                         device=logits.device)
+    spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
 
+    mock_sampler_output(rejection_sampler, bonus_token_tensor)
     output = rejection_sampler(
         spec_decode_metadata,
         draft_probs=None,
-        target_logits=logits,
-        bonus_token_ids=bonus_token_tensor,
+        logits=logits,
         sampling_metadata=metadata,
     )
     expected = torch.tensor([[5]], dtype=torch.int, device=logits.device)
-    assert torch.equal(output, expected)
+    assert torch.equal(output.sampled_token_ids, expected)
 
 
 def test_multiple_mismatches(rejection_sampler):
@@ -208,14 +222,13 @@ def test_multiple_mismatches(rejection_sampler):
     logits = create_logits_tensor(output_tokens)
     bonus_token_tensor = torch.tensor(
         [output_tokens[0][-1], output_tokens[1][-1]], device=logits.device)
-    spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
-                                                         device=logits.device)
+    spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
 
+    mock_sampler_output(rejection_sampler, bonus_token_tensor)
     output = rejection_sampler(
         spec_decode_metadata,
         draft_probs=None,
-        target_logits=logits,
-        bonus_token_ids=bonus_token_tensor,
+        logits=logits,
         sampling_metadata=metadata,
     )
     expected = torch.tensor(
@@ -224,7 +237,7 @@ def test_multiple_mismatches(rejection_sampler):
         dtype=torch.int,
         device=logits.device,
     )
-    assert torch.equal(output, expected)
+    assert torch.equal(output.sampled_token_ids, expected)
 
 
 @pytest.mark.parametrize(
@@ -242,20 +255,19 @@ def test_parametrized_cases(rejection_sampler, spec_tokens, output_tokens,
     logits = create_logits_tensor(output_tokens)
     bonus_token_tensor = torch.tensor([tokens[-1] for tokens in output_tokens],
                                       device=logits.device)
-    spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens,
-                                                         device=logits.device)
+    spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
 
+    mock_sampler_output(rejection_sampler, bonus_token_tensor)
     output = rejection_sampler(
         spec_decode_metadata,
         draft_probs=None,
-        target_logits=logits,
-        bonus_token_ids=bonus_token_tensor,
+        logits=logits,
         sampling_metadata=metadata,
     )
     expected_tensor = torch.tensor(expected,
                                    dtype=torch.int,
                                    device=logits.device)
-    assert torch.equal(output, expected_tensor)
+    assert torch.equal(output.sampled_token_ids, expected_tensor)
 
 
 ########################### Tests for Random Sampling ###################
@@ -305,17 +317,18 @@ def test_deterministic_when_seeded(
         sampling_metadata = create_sampling_metadata(all_greedy=False,
                                                      temperature=temperature,
                                                      generators=seeded_seqs)
-        spec_decode_metadata = SpecDecodeMetadata.make_dummy(
-            draft_token_ids.tolist(), device=DEVICE)
+        spec_decode_metadata = create_spec_decode_metadata(
+            draft_token_ids.tolist(), target_logits)
+
+        mock_sampler_output(rejection_sampler, bonus_token_ids)
         rep_result = rejection_sampler(
             spec_decode_metadata,
-            draft_probs=draft_probs,
-            target_logits=target_logits,
-            bonus_token_ids=bonus_token_ids,
+            draft_probs=None,
+            logits=target_logits,
             sampling_metadata=sampling_metadata,
         )
 
-        results.append(rep_result)
+        results.append(rep_result.sampled_token_ids)
 
     for i in range(batch_size):
         if seeded_mask[i]:
@@ -424,7 +437,9 @@ def estimate_rejection_sampling_pdf(
     Returns:
         Estimated probability distribution of the output tokens.
     """
-    rejection_sampler = RejectionSampler()
+    # Mock the sampler that TreeRejectionSampler uses
+    mock_sampler = Mock(spec=Sampler)
+    rejection_sampler = RejectionSampler(mock_sampler)
     num_tokens = num_samples * k
     # Repeat draft probs num_samples * k times.
     draft_probs = draft_probs.reshape(1, 1,
@@ -447,16 +462,17 @@ def estimate_rejection_sampling_pdf(
     temperature = torch.ones(num_samples, dtype=torch.float32, device=DEVICE)
     sampling_metadata = create_sampling_metadata(all_greedy=False,
                                                  temperature=temperature)
-    spec_decode_metadata = SpecDecodeMetadata.make_dummy(
-        draft_token_ids.tolist(), device=bonus_token_ids.device)
-    output_token_ids = rejection_sampler(
+    spec_decode_metadata = create_spec_decode_metadata(
+        draft_token_ids.tolist(), target_logits)
+
+    mock_sampler_output(rejection_sampler, bonus_token_ids)
+    sampler_output = rejection_sampler(
         spec_decode_metadata,
         draft_probs=draft_probs,
-        target_logits=target_logits,
-        bonus_token_ids=bonus_token_ids,
+        logits=target_logits,
         sampling_metadata=sampling_metadata,
     )
-    output_token_ids = output_token_ids[:, :-1].flatten()
+    output_token_ids = sampler_output.sampled_token_ids[:, :-1].flatten()
 
     hist = torch.histogram(output_token_ids.to(dtype=torch.float,
                                                device="cpu"),
@@ -496,22 +512,20 @@ def _test_masked_logits(
                                   device=DEVICE)
 
     # Create spec decode metadata
-    spec_decode_metadata = SpecDecodeMetadata.make_dummy(
-        draft_token_ids,
-        device=DEVICE,
-    )
+    spec_decode_metadata = create_spec_decode_metadata(draft_token_ids,
+                                                       target_logits)
 
     # Run rejection sampling
-    output_token_ids = rejection_sampler(
+    mock_sampler_output(rejection_sampler, bonus_token_ids)
+    output = rejection_sampler(
         spec_decode_metadata,
         draft_probs=draft_probs,
-        target_logits=target_logits,
-        bonus_token_ids=bonus_token_ids,
+        logits=target_logits,
         sampling_metadata=sampling_metadata,
     )
 
     # Remove bonus tokens and reshape
-    output_token_ids = output_token_ids[:, :-1].flatten().tolist()
+    output_token_ids = output.sampled_token_ids[:, :-1].flatten().tolist()
 
     # Check that all sampled tokens are within the unmasked indices.
     for i in range(num_tokens):

vllm/v1/engine/logprobs.py

@@ -66,7 +66,7 @@ def _update_sample_logprobs(self, logprobs_lists: LogprobsLists) -> None:
         assert self.logprobs is not None
         assert self.cumulative_logprob is not None
 
-        token_ids_lst, logprobs_lst, ranks_lst = logprobs_lists
+        token_ids_lst, logprobs_lst, ranks_lst, _ = logprobs_lists
 
         for rank, logprobs, token_ids in zip(ranks_lst, logprobs_lst,
                                              token_ids_lst):