[sgl] add ability to return logprobs in MultiLayerEagleWorkerV2 (#22241)

python/sglang/srt/layers/utils/logprob.py

@@ -427,3 +427,68 @@ def add_output_logprobs_for_spec_v1(
                     req.output_token_ids_logprobs_val.append(token_ids_logprobs_val[pt])
                     req.output_token_ids_logprobs_idx.append(token_ids_logprobs_idx[pt])
             pt += 1
+
+
+def compute_spec_v2_logprobs(
+    batch,
+    logits_output,
+    predict: torch.Tensor,
+    accept_index: torch.Tensor,
+    speculative_num_steps: int,
+):
+    """Compute logprobs for accepted tokens after spec v2 verify sampling.
+
+    Gathers logits at accepted positions, applies log_softmax (temperature-scaled
+    if not greedy), and populates logits_output.next_token_logprobs (plus optional
+    top-k / token-ids logprobs) so they flow through copy_to_cpu().
+    """
+    bs = len(batch.seq_lens)
+    max_accept = speculative_num_steps + 1
+    device = predict.device
+
+    flat_accept_idx = accept_index.long().reshape(-1)
+    gathered_logits = logits_output.next_token_logits[flat_accept_idx]
+
+    if batch.sampling_info.is_all_greedy or envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get():
+        gathered_logprobs = torch.nn.functional.log_softmax(gathered_logits, dim=-1)
+    else:
+        temperatures = torch.repeat_interleave(
+            batch.sampling_info.temperatures,
+            max_accept,
+            dim=0,
+        )
+        gathered_logprobs = torch.nn.functional.log_softmax(
+            gathered_logits / temperatures, dim=-1
+        )
+    gathered_logprobs.clamp_(min=torch.finfo(gathered_logprobs.dtype).min)
+
+    accepted_token_ids = predict[flat_accept_idx]
+    token_logprobs = gathered_logprobs[
+        torch.arange(bs * max_accept, device=device),
+        accepted_token_ids.long(),
+    ]
+    logits_output.next_token_logprobs = token_logprobs.reshape(bs, max_accept)
+
+    if batch.top_logprobs_nums and any(x > 0 for x in batch.top_logprobs_nums):
+        top_logprobs_nums_expanded = [
+            num for num in batch.top_logprobs_nums for _ in range(max_accept)
+        ]
+        (
+            logits_output.next_token_top_logprobs_val,
+            logits_output.next_token_top_logprobs_idx,
+        ) = get_top_logprobs(
+            gathered_logprobs, top_logprobs_nums_expanded, no_copy_to_cpu=True
+        )
+
+    if batch.token_ids_logprobs and any(
+        x is not None for x in batch.token_ids_logprobs
+    ):
+        token_ids_logprobs_expanded = [
+            ids for ids in batch.token_ids_logprobs for _ in range(max_accept)
+        ]
+        (
+            logits_output.next_token_token_ids_logprobs_val,
+            logits_output.next_token_token_ids_logprobs_idx,
+        ) = get_token_ids_logprobs(
+            gathered_logprobs, token_ids_logprobs_expanded, no_copy_to_cpu=True
+        )

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -17,12 +17,11 @@ from sglang.srt.layers.attention.trtllm_mla_backend import (
     TRTLLMMLABackend,
 )
 from sglang.srt.layers.dp_attention import get_attention_tp_group
-from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.moe.utils import (
     speculative_moe_a2a_backend_context,
     speculative_moe_backend_context,
 )
-from sglang.srt.layers.utils.logprob import get_token_ids_logprobs, get_top_logprobs
+from sglang.srt.layers.utils.logprob import compute_spec_v2_logprobs
 from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput
 from sglang.srt.managers.schedule_batch import ModelWorkerBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
@@ -852,7 +851,9 @@ class EAGLEWorkerV2(BaseSpecWorker):
             verified_id = torch.empty((0,), device=self.device, dtype=torch.int32)
 
         if batch.return_logprob and not batch.forward_mode.is_idle():
-            self._compute_spec_v2_logprobs(batch, logits_output, predict, accept_index)
+            compute_spec_v2_logprobs(
+                batch, logits_output, predict, accept_index, self.speculative_num_steps
+            )
 
         # Construct the next draft input
         next_draft_input = EagleDraftInput(
@@ -869,72 +870,6 @@ class EAGLEWorkerV2(BaseSpecWorker):
             accept_lens=accept_length,
         )
 
-    def _compute_spec_v2_logprobs(
-        self,
-        batch: ModelWorkerBatch,
-        logits_output: LogitsProcessorOutput,
-        predict: torch.Tensor,
-        accept_index: torch.Tensor,
-    ):
-        """Compute logprobs for accepted tokens on GPU in the forward stream.
-
-        Stores results in logits_output fields so they flow through copy_to_cpu().
-        """
-
-        bs = len(batch.seq_lens)
-        max_accept = self.speculative_num_steps + 1
-        device = predict.device
-
-        flat_accept_idx = accept_index.long().reshape(-1)
-        gathered_logits = logits_output.next_token_logits[flat_accept_idx]
-
-        if (
-            batch.sampling_info.is_all_greedy
-            or envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get()
-        ):
-            gathered_logprobs = torch.nn.functional.log_softmax(gathered_logits, dim=-1)
-        else:
-            temperatures = torch.repeat_interleave(
-                batch.sampling_info.temperatures,
-                max_accept,
-                dim=0,
-            )
-            gathered_logprobs = torch.nn.functional.log_softmax(
-                gathered_logits / temperatures, dim=-1
-            )
-        gathered_logprobs.clamp_(min=torch.finfo(gathered_logprobs.dtype).min)
-
-        accepted_token_ids = predict[flat_accept_idx]
-        token_logprobs = gathered_logprobs[
-            torch.arange(bs * max_accept, device=device),
-            accepted_token_ids.long(),
-        ]
-        logits_output.next_token_logprobs = token_logprobs.reshape(bs, max_accept)
-
-        if batch.top_logprobs_nums and any(x > 0 for x in batch.top_logprobs_nums):
-            top_logprobs_nums_expanded = [
-                num for num in batch.top_logprobs_nums for _ in range(max_accept)
-            ]
-            (
-                logits_output.next_token_top_logprobs_val,
-                logits_output.next_token_top_logprobs_idx,
-            ) = get_top_logprobs(
-                gathered_logprobs, top_logprobs_nums_expanded, no_copy_to_cpu=True
-            )
-
-        if batch.token_ids_logprobs and any(
-            x is not None for x in batch.token_ids_logprobs
-        ):
-            token_ids_logprobs_expanded = [
-                ids for ids in batch.token_ids_logprobs for _ in range(max_accept)
-            ]
-            (
-                logits_output.next_token_token_ids_logprobs_val,
-                logits_output.next_token_token_ids_logprobs_idx,
-            ) = get_token_ids_logprobs(
-                gathered_logprobs, token_ids_logprobs_expanded, no_copy_to_cpu=True
-            )
-
     def _mamba_verify_update(
         self,
         batch: ModelWorkerBatch,

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -20,6 +20,7 @@ import torch
 
 from sglang.srt.environ import envs
 from sglang.srt.layers.moe.utils import speculative_moe_backend_context
+from sglang.srt.layers.utils.logprob import compute_spec_v2_logprobs
 from sglang.srt.managers.schedule_batch import ModelWorkerBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.managers.tp_worker import TpModelWorker
@@ -746,6 +747,11 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
         else:
             verified_id = torch.empty((0,), device=self.device, dtype=torch.int32)
 
+        if batch.return_logprob and not batch.forward_mode.is_idle():
+            compute_spec_v2_logprobs(
+                batch, logits_output, predict, accept_index, self.speculative_num_steps
+            )
+
         # Construct the next draft input
         next_draft_input = EagleDraftInput(
             verified_id=verified_id,