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

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,

142
test/registered/spec/eagle/test_eagle_infer_beta.py
View File

@@ -1,6 +1,9 @@
import unittest
from types import SimpleNamespace
import numpy as np
import requests
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
@@ -98,6 +101,145 @@ class TestEagleServerBase(CustomTestCase, MatchedStopMixin):
) # 0.3333 for 60 questions; 0.234 for 1319 questions
assert self.process.poll() is None
def test_logprob_spec_v2_match(self):
"""Verify spec v2 decode logprobs match prefill scoring logprobs.
Generate tokens with spec v2, then score the same sequence via
prefill-only (no speculation). The two sets of logprobs should be
close, validating that spec v2 computes logprobs correctly.
Runs two rounds with different prompts to catch state-dependent bugs.
"""
top_k = 5
probe_token_ids = [1, 2, 10, 100, 1000]
prompts = [
"The capital of France is",
"Explain quantum computing in simple terms:",
]
for round_idx, prompt in enumerate(prompts):
with self.subTest(round=round_idx, prompt=prompt):
gen_res = requests.post(
self.base_url + "/generate",
json={
"text": prompt,
"sampling_params": {
"temperature": 0,
"max_new_tokens": 32,
"ignore_eos": True,
},
"return_logprob": True,
"top_logprobs_num": top_k,
"token_ids_logprob": probe_token_ids,
"logprob_start_len": 0,
},
).json()
decode_logprobs = gen_res["meta_info"]["output_token_logprobs"]
decode_top_logprobs = gen_res["meta_info"]["output_top_logprobs"]
decode_tid_logprobs = gen_res["meta_info"]["output_token_ids_logprobs"]
input_token_ids = [
t[1] for t in gen_res["meta_info"]["input_token_logprobs"]
]
output_token_ids = [t[1] for t in decode_logprobs]
num_prompt_tokens = gen_res["meta_info"]["prompt_tokens"]
score_res = requests.post(
self.base_url + "/generate",
json={
"input_ids": input_token_ids + output_token_ids,
"sampling_params": {
"temperature": 0,
"max_new_tokens": 0,
},
"return_logprob": True,
"top_logprobs_num": top_k,
"token_ids_logprob": probe_token_ids,
"logprob_start_len": 0,
},
).json()
score_logprobs = score_res["meta_info"]["input_token_logprobs"][
num_prompt_tokens:
]
score_top_logprobs = score_res["meta_info"]["input_top_logprobs"][
num_prompt_tokens:
]
score_tid_logprobs = score_res["meta_info"]["input_token_ids_logprobs"][
num_prompt_tokens:
]
self.assertEqual(len(decode_logprobs), len(score_logprobs))
# Check per-token logprobs
decode_vals = np.array([t[0] for t in decode_logprobs])
score_vals = np.array([t[0] for t in score_logprobs])
max_diff = np.max(np.abs(decode_vals - score_vals))
print(
f"[round {round_idx}] prompt={prompt!r} "
f"logprob max_diff={max_diff:.6f}"
)
print(f"[round {round_idx}] decode_vals[-5:]={decode_vals[-5:]}")
print(f"[round {round_idx}] score_vals[-5:]={score_vals[-5:]}")
self.assertLess(max_diff, 0.255)
# Check top-k logprobs
for pos in range(len(decode_logprobs)):
dec_top = {t[1]: t[0] for t in decode_top_logprobs[pos]}
scr_top = {t[1]: t[0] for t in score_top_logprobs[pos]}
common_ids = set(dec_top.keys()) & set(scr_top.keys())
self.assertGreater(len(common_ids), 0)
for tid in common_ids:
self.assertAlmostEqual(dec_top[tid], scr_top[tid], delta=0.255)
# Check token_ids_logprob
self.assertEqual(len(decode_tid_logprobs), len(score_tid_logprobs))
for pos in range(len(decode_tid_logprobs)):
dec_tid = {t[1]: t[0] for t in decode_tid_logprobs[pos]}
scr_tid = {t[1]: t[0] for t in score_tid_logprobs[pos]}
self.assertEqual(set(dec_tid.keys()), set(scr_tid.keys()))
for tid in dec_tid:
self.assertAlmostEqual(dec_tid[tid], scr_tid[tid], delta=0.255)
def test_token_ids_logprob_ragged(self):
"""Regression: get_token_ids_logprobs_raw crashes on ragged token_ids_logprob lists.
Sends concurrent requests with different-length token_ids_logprob lists
so they land in the same batch. torch.tensor() on ragged input will crash.
"""
import concurrent.futures
def send(probe_ids):
return requests.post(
self.base_url + "/generate",
json={
"text": "Hello world",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 8,
},
"return_logprob": True,
"top_logprobs_num": 3,
"token_ids_logprob": probe_ids,
},
).json()
ragged_probes = [
[1, 2],
[3, 4, 5],
[6],
[10, 20, 30, 40],
[1, 2],
[3, 4, 5],
[6],
[10, 20, 30, 40],
]
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
futs = [pool.submit(send, ids) for ids in ragged_probes]
for f in concurrent.futures.as_completed(futs):
res = f.result()
self.assertIn("text", res, f"Server error: {res}")
class TestEagleServerPage(TestEagleServerBase):
other_launch_args = ["--page-size", "64"]