chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support return_logprob for spec v2 (overlap safe) (#19801)

Browse Source 
Co-authored-by: Ratish1 <ratish1501@gmail.com>
Co-authored-by: Ratish1 <formula733@gmail.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
This commit is contained in:
Qiaolin Yu
2026-03-10 15:38:27 -07:00
committed by GitHub
parent 76ee4bb98c
commit 09a118fafe
6 changed files with 314 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
python/sglang/srt/speculative/eagle_worker_v2.py
View File

@@ -17,10 +17,12 @@ from sglang.srt.layers.attention.trtllm_mla_backend import (
TRTLLMMLAMultiStepDraftBackend,
)
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.managers.io_struct import UpdateWeightsFromTensorReqInput
from sglang.srt.managers.schedule_batch import ModelWorkerBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
@@ -834,6 +836,9 @@ class EAGLEWorkerV2(BaseSpecWorker):
else:
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)
# Construct the next draft input
next_draft_input = EagleDraftInput(
verified_id=verified_id,
@@ -849,6 +854,72 @@ 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,