From cebd9c2a1e00fcdd50f420f2a685ad4a26cb39c3 Mon Sep 17 00:00:00 2001 From: Tarushii Goel Date: Thu, 9 Apr 2026 16:20:55 -0700 Subject: [PATCH] [sgl] add ability to return logprobs in MultiLayerEagleWorkerV2 (#22241) --- python/sglang/srt/layers/utils/logprob.py | 65 +++++++++++++++++ .../sglang/srt/speculative/eagle_worker_v2.py | 73 +------------------ .../multi_layer_eagle_worker_v2.py | 6 ++ 3 files changed, 75 insertions(+), 69 deletions(-) diff --git a/python/sglang/srt/layers/utils/logprob.py b/python/sglang/srt/layers/utils/logprob.py index e7ce243fe..0bdb0c87f 100644 --- a/python/sglang/srt/layers/utils/logprob.py +++ b/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 + ) diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 86f72527f..b15196804 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 65bfb048e..a0829da9d 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/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,