chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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

2
python/sglang/srt/layers/logits_processor.py
View File

@@ -1055,7 +1055,7 @@ class LogitsProcessor(nn.Module):
input_token_ids_logprobs_val,
input_token_ids_logprobs_idx,
) = get_token_ids_logprobs_prefill(
sliced_logprobs, logits_metadata, delay_cpu_copy=True
sliced_logprobs, logits_metadata, no_copy_to_cpu=True
)
# Get the logprob of top-k tokens

48
python/sglang/srt/layers/utils/logprob.py
View File

@@ -68,11 +68,13 @@ def get_top_logprobs_raw(
top_logprobs_nums: List[int],
stage: LogprobStage,
extend_logprob_pruned_lens_cpu: Optional[List[int]] = None,
no_copy_to_cpu: bool = False,
):
max_k = max(top_logprobs_nums)
values, indices = logprobs.topk(max_k, dim=-1)
values = values.tolist()
indices = indices.tolist()
if not no_copy_to_cpu:
values = values.tolist()
indices = indices.tolist()
top_logprobs_val = []
top_logprobs_idx = []
@@ -110,57 +112,73 @@ def get_top_logprobs_prefill(
def get_top_logprobs(
logprobs: torch.Tensor,
top_logprobs_nums: List[int],
no_copy_to_cpu: bool = False,
):
return get_top_logprobs_raw(logprobs, top_logprobs_nums, stage=LogprobStage.DECODE)
return get_top_logprobs_raw(
logprobs,
top_logprobs_nums,
stage=LogprobStage.DECODE,
no_copy_to_cpu=no_copy_to_cpu,
)
def get_token_ids_logprobs_raw(
logprobs: torch.Tensor,
token_ids_logprobs: List[Optional[List[int]]],
token_ids_logprobs_list: List[Optional[List[int]]],
stage: LogprobStage,
extend_logprob_pruned_lens_cpu: Optional[List[int]] = None,
delay_cpu_copy: bool = False,
no_copy_to_cpu: bool = False,
):
vals, idxs = [], []
if stage == LogprobStage.DECODE:
for i, token_ids in enumerate(token_ids_logprobs):
for i, token_ids in enumerate(token_ids_logprobs_list):
if token_ids is None:
vals.append([])
idxs.append([])
else:
vals.append(logprobs[i, token_ids].tolist())
token_ids_tensor = torch.tensor(token_ids, dtype=torch.long).to(
logprobs.device, non_blocking=True
)
row = logprobs[i, token_ids_tensor]
vals.append(row if no_copy_to_cpu else row.tolist())
idxs.append(token_ids)
else: # prefill
pt = 0
for token_ids, pruned_len in zip(
token_ids_logprobs, extend_logprob_pruned_lens_cpu
for i, (token_ids, pruned_len) in enumerate(
zip(token_ids_logprobs_list, extend_logprob_pruned_lens_cpu)
):
if pruned_len <= 0:
vals.append([])
idxs.append([])
continue
pos_logprobs = logprobs[pt : pt + pruned_len, token_ids]
vals.append(pos_logprobs if delay_cpu_copy else pos_logprobs.tolist())
token_ids_tensor = torch.tensor(token_ids, dtype=torch.long).to(
logprobs.device, non_blocking=True
)
pos_logprobs = logprobs[pt : pt + pruned_len, token_ids_tensor]
vals.append(pos_logprobs if no_copy_to_cpu else pos_logprobs.tolist())
idxs.append([token_ids for _ in range(pruned_len)])
pt += pruned_len
return vals, idxs
def get_token_ids_logprobs_prefill(
all_logprobs, logits_metadata: LogitsMetadata, delay_cpu_copy=False
all_logprobs, logits_metadata: LogitsMetadata, no_copy_to_cpu=False
):
return get_token_ids_logprobs_raw(
all_logprobs,
logits_metadata.token_ids_logprobs,
stage=LogprobStage.PREFILL,
extend_logprob_pruned_lens_cpu=logits_metadata.extend_logprob_pruned_lens_cpu,
delay_cpu_copy=delay_cpu_copy,
no_copy_to_cpu=no_copy_to_cpu,
)
def get_token_ids_logprobs(logprobs, token_ids_logprobs):
def get_token_ids_logprobs(logprobs, token_ids_logprobs, no_copy_to_cpu=False):
return get_token_ids_logprobs_raw(
logprobs, token_ids_logprobs, stage=LogprobStage.DECODE
logprobs,
token_ids_logprobs,
stage=LogprobStage.DECODE,
no_copy_to_cpu=no_copy_to_cpu,
)

74
python/sglang/srt/managers/scheduler_output_processor_mixin.py
View File

@@ -367,12 +367,27 @@ class SchedulerOutputProcessorMixin:
result.can_run_cuda_graph,
)
if batch.spec_algorithm.is_none():
next_token_ids = next_token_ids.tolist()
if batch.spec_algorithm.is_none() or batch.is_spec_v2:
if batch.is_spec_v2:
next_token_ids = self._resolve_spec_overlap_token_ids(result, batch)
else:
next_token_ids = next_token_ids.tolist()
if batch.return_logprob:
next_token_logprobs = logits_output.next_token_logprobs.tolist()
elif batch.is_spec_v2:
next_token_ids = self._resolve_spec_overlap_token_ids(result, batch)
if batch.is_spec_v2 and logits_output.next_token_top_logprobs_val:
logits_output.next_token_top_logprobs_val = [
v.tolist() for v in logits_output.next_token_top_logprobs_val
]
logits_output.next_token_top_logprobs_idx = [
x.tolist() for x in logits_output.next_token_top_logprobs_idx
]
if batch.is_spec_v2 and logits_output.next_token_token_ids_logprobs_val:
logits_output.next_token_token_ids_logprobs_val = [
v.tolist()
for v in logits_output.next_token_token_ids_logprobs_val
]
self.num_generated_tokens += len(batch.reqs)
if not batch.spec_algorithm.is_none():
@@ -439,24 +454,39 @@ class SchedulerOutputProcessorMixin:
self.maybe_collect_customized_info(i, req, logits_output)
if req.return_logprob and batch.spec_algorithm.is_none():
# speculative worker handles logprob in speculative decoding
req.output_token_logprobs_val.append(next_token_logprobs[i])
req.output_token_logprobs_idx.append(next_token_id)
if req.top_logprobs_num > 0:
req.output_top_logprobs_val.append(
logits_output.next_token_top_logprobs_val[i]
)
req.output_top_logprobs_idx.append(
logits_output.next_token_top_logprobs_idx[i]
)
if req.token_ids_logprob is not None:
req.output_token_ids_logprobs_val.append(
logits_output.next_token_token_ids_logprobs_val[i]
)
req.output_token_ids_logprobs_idx.append(
logits_output.next_token_token_ids_logprobs_idx[i]
)
if req.return_logprob and (
batch.spec_algorithm.is_none() or batch.is_spec_v2
):
# Spec v1 handles logprobs inside its own worker.
# Normalize: non-spec has 1 token, spec v2 has multiple.
if batch.is_spec_v2:
accepted_logprobs = next_token_logprobs[i]
accepted_ids = next_token_id
max_accept = len(accepted_logprobs)
else:
accepted_logprobs = [next_token_logprobs[i]]
accepted_ids = [next_token_id]
max_accept = 1
for j, tok_id in enumerate(accepted_ids):
req.output_token_logprobs_val.append(accepted_logprobs[j])
req.output_token_logprobs_idx.append(tok_id)
if req.top_logprobs_num > 0:
flat_idx = i * max_accept + j
req.output_top_logprobs_val.append(
logits_output.next_token_top_logprobs_val[flat_idx]
)
req.output_top_logprobs_idx.append(
logits_output.next_token_top_logprobs_idx[flat_idx]
)
if req.token_ids_logprob is not None:
flat_idx = i * max_accept + j
req.output_token_ids_logprobs_val.append(
logits_output.next_token_token_ids_logprobs_val[flat_idx]
)
req.output_token_ids_logprobs_idx.append(
logits_output.next_token_token_ids_logprobs_idx[flat_idx]
)
if req.return_hidden_states and logits_output.hidden_states is not None:
req.hidden_states.append(

15
python/sglang/srt/managers/utils.py
View File

@@ -63,6 +63,21 @@ class GenerationBatchResult:
self.logits_output.input_token_logprobs = (
self.logits_output.input_token_logprobs.to("cpu", non_blocking=True)
)
if self.logits_output.next_token_top_logprobs_val is not None:
self.logits_output.next_token_top_logprobs_val = [
v.to("cpu", non_blocking=True) if torch.is_tensor(v) else v
for v in self.logits_output.next_token_top_logprobs_val
]
if self.logits_output.next_token_top_logprobs_idx is not None:
self.logits_output.next_token_top_logprobs_idx = [
x.to("cpu", non_blocking=True) if torch.is_tensor(x) else x
for x in self.logits_output.next_token_top_logprobs_idx
]
if self.logits_output.next_token_token_ids_logprobs_val is not None:
self.logits_output.next_token_token_ids_logprobs_val = [
v.to("cpu", non_blocking=True) if torch.is_tensor(v) else v
for v in self.logits_output.next_token_token_ids_logprobs_val
]
if self.logits_output.hidden_states is not None:
self.logits_output.hidden_states = self.logits_output.hidden_states.to(
"cpu", non_blocking=True

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,