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

Signed-off-by: Giancarlo Delfin <gdelfin@meta.com>

Author: TheEpicDolphin

tests/v1/sample/test_logprobs.py

@@ -9,6 +9,7 @@
 import pytest
 import torch
 
+from tests.utils import large_gpu_mark
 from tests.v1.sample.utils import (
     BatchLogprobsComposition,
     BatchLogprobsSpecType,
@@ -18,6 +19,7 @@
 )
 from vllm import SamplingParams
 from vllm.config.model import LogprobsMode
+from vllm.distributed import cleanup_dist_env_and_memory
 
 from ...conftest import HfRunner, VllmRunner
 
@@ -515,11 +517,14 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode):
 @pytest.mark.parametrize(
     "model_setup",
     [
-        (
-            "eagle",
-            "meta-llama/Llama-3.1-8B-Instruct",
-            "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
-        )
+        pytest.param(
+            (
+                "eagle",
+                "meta-llama/Llama-3.1-8B-Instruct",
+                "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
+            ),
+            marks=large_gpu_mark(min_gb=32),
+        ),
     ],
 )
 def test_spec_decode_logprobs(
@@ -543,17 +548,27 @@ def test_spec_decode_logprobs(
             temperature=0, logprobs=3, max_tokens=10, ignore_eos=False
         )
         method, model_name, spec_model_name = model_setup
+        max_model_len = 256
 
         # Run base LLM.
         ref_llm = LLM(
             model=model_name,
             max_logprobs=5,
-            max_model_len=2048,
+            max_model_len=max_model_len,
             seed=42,
             logprobs_mode=logprobs_mode,
+            gpu_memory_utilization=0.4,
         )
         ref_results = ref_llm.generate([prompt], sampling_params)
+        # Collect logprobs outputs from reference LLM.
+        ref_logprobs = []
+        for output in ref_results[0].outputs:
+            for logprobs in output.logprobs:
+                for token_id in logprobs:
+                    ref_logprobs.append(logprobs[token_id])
         del ref_llm
+        torch.cuda.empty_cache()
+        cleanup_dist_env_and_memory()
 
         # Run spec decode LLM.
         spec_llm = LLM(
@@ -562,27 +577,24 @@ def test_spec_decode_logprobs(
                 "method": method,
                 "model": spec_model_name,
                 "num_speculative_tokens": 3,
-                "max_model_len": 2048,
+                "max_model_len": max_model_len,
             },
             max_logprobs=5,
-            max_model_len=2048,
+            max_model_len=max_model_len,
             seed=42,
             logprobs_mode=logprobs_mode,
+            gpu_memory_utilization=0.4,
         )
         spec_results = spec_llm.generate([prompt], sampling_params)
-        del spec_llm
-
-        # Collect logprobs outputs from reference and spec decode LLMs.
-        ref_logprobs = []
-        for output in ref_results[0].outputs:
-            for logprobs in output.logprobs:
-                for token_id in logprobs:
-                    ref_logprobs.append(logprobs[token_id])
+        # Collect logprobs outputs from spec decode LLM.
         spec_logprobs = []
         for output in spec_results[0].outputs:
             for logprobs in output.logprobs:
                 for token_id in logprobs:
                     spec_logprobs.append(logprobs[token_id])
+        del spec_llm
+        torch.cuda.empty_cache()
+        cleanup_dist_env_and_memory()
 
         # Per-token logprobs are expected to be the same.
         assert len(ref_logprobs) == len(spec_logprobs)

tests/v1/sample/test_rejection_sampler.py

@@ -21,6 +21,7 @@
 @pytest.fixture
 def rejection_sampler():
     mock_sampler = Mock(spec=Sampler)
+    mock_sampler.logprobs_mode = "raw_logprobs"
     return RejectionSampler(mock_sampler)
 
 
@@ -469,8 +470,8 @@ def estimate_rejection_sampling_pdf(
     Returns:
         Estimated probability distribution of the output tokens.
     """
-    # Mock the sampler that TreeRejectionSampler uses
     mock_sampler = Mock(spec=Sampler)
+    mock_sampler.logprobs_mode = "raw_logprobs"
     rejection_sampler = RejectionSampler(mock_sampler)
     num_tokens = num_samples * k
     # Repeat draft probs num_samples * k times.
@@ -674,19 +675,19 @@ def test_frequency_penalties(rejection_sampler):
     spec_decode_metadata = SpecDecodeMetadata.make_dummy(
         spec_tokens, device=logits.device
     )
+    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, 15, -1, -1], [7, -1, -1, -1], [1, 1, 15, -1]],
         dtype=torch.int,
         device=logits.device,
     )
-    assert torch.equal(output, expected)
+    assert torch.equal(output.sampled_token_ids, expected)
 
 
 def test_bad_words(rejection_sampler):
@@ -716,14 +717,12 @@ def test_bad_words(rejection_sampler):
     bonus_token_tensor = torch.tensor(
         [output_tokens[i][-1] for i in range(len(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,
     )
 
@@ -732,7 +731,7 @@ def test_bad_words(rejection_sampler):
         dtype=torch.int,
         device=logits.device,
     )
-    assert torch.equal(output, expected)
+    assert torch.equal(output.sampled_token_ids, expected)
 
 
 def test_allowed_token_ids(rejection_sampler):
@@ -765,14 +764,12 @@ def test_allowed_token_ids(rejection_sampler):
     bonus_token_tensor = torch.tensor(
         [output_tokens[i][-1] for i in range(len(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,
     )
 
@@ -781,4 +778,4 @@ def test_allowed_token_ids(rejection_sampler):
         dtype=torch.int,
         device=logits.device,
     )
-    assert torch.equal(output, expected)
+    assert torch.equal(output.sampled_token_ids, expected)

vllm/v1/sample/rejection_sampler.py

@@ -51,9 +51,9 @@ class RejectionSampler(nn.Module):
     def __init__(self, sampler: Sampler):
         super().__init__()
         self.sampler = sampler
-        self.return_processed_logprobs = self.sampler.logprobs_mode.startswith(
-            "processed"
-        )
+        logprobs_mode = self.sampler.logprobs_mode
+        self.is_processed_logprobs_mode = logprobs_mode.startswith("processed")
+        self.is_logits_logprobs_mode = logprobs_mode.endswith("logits")
 
     def forward(
         self,
@@ -107,10 +107,9 @@ def forward(
             # Override the logprobs mode to return logits because they are
             # needed later to compute the accepted token logprobs.
             logprobs_mode_override="processed_logits"
-            if self.return_processed_logprobs
+            if self.is_processed_logprobs_mode
             else "raw_logits",
         )
-        bonus_logits = bonus_sampler_output.logprobs_tensors.logprobs
         bonus_token_ids = bonus_sampler_output.sampled_token_ids
 
         # Just like `bonus_logits`, `target_logits` is a new tensor with
@@ -144,16 +143,21 @@ def forward(
             sampling_metadata,
         )
 
-        return SamplerOutput(
-            sampled_token_ids=output_token_ids,
-            logprobs_tensors=self._get_logprobs_tensors(
+        logprobs_tensors = None
+        if sampling_metadata.max_num_logprobs:
+            bonus_logits = bonus_sampler_output.logprobs_tensors.logprobs
+            logprobs_tensors = self._get_logprobs_tensors(
                 sampling_metadata,
                 metadata,
                 logits,
-                target_logits if self.return_processed_logprobs else raw_target_logits,
+                target_logits if self.is_processed_logprobs_mode else raw_target_logits,
                 bonus_logits,
                 output_token_ids,
-            ),
+            )
+
+        return SamplerOutput(
+            sampled_token_ids=output_token_ids,
+            logprobs_tensors=logprobs_tensors,
         )
 
     def _get_logprobs_tensors(
@@ -164,10 +168,7 @@ def _get_logprobs_tensors(
         target_logits: torch.Tensor,
         bonus_logits: torch.Tensor,
         sampled_token_ids: torch.Tensor,
-    ) -> LogprobsTensors | None:
-        if not sampling_metadata.max_num_logprobs:
-            return None
-
+    ) -> LogprobsTensors:
         cu_num_sampled_tokens = torch.zeros_like(metadata.cu_num_sampled_tokens)
         cu_num_sampled_tokens[1:] = metadata.cu_num_sampled_tokens[:-1]
 
@@ -190,7 +191,7 @@ def _get_logprobs_tensors(
         accepted_logits = final_logits[accepted_logit_indices]
         accepted_logprobs = (
             accepted_logits
-            if self.logprobs_mode.endswith("logits")
+            if self.is_logits_logprobs_mode
             else self.sampler.compute_logprobs(accepted_logits)
         )
         accepted_tokens = sampled_token_ids[accepted_mask]