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

[2/N] clean duplicate code of logprob processing in spec. (#15593)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-12-24 15:21:13 +08:00
committed by GitHub
parent ff903a7eea
commit cb719c74ad
5 changed files with 127 additions and 267 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2,12 +2,17 @@ from __future__ import annotations
import dataclasses
from enum import Enum, auto
from typing import TYPE_CHECKING, List, Optional
from typing import TYPE_CHECKING, List, Optional, Union
import torch
from sglang.srt.environ import envs
if TYPE_CHECKING:
from sglang.srt.layers.logits_processor import LogitsMetadata
from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.speculative.eagle_info import EagleVerifyOutput
from sglang.srt.speculative.ngram_info import NgramVerifyInput
class LogprobStage(Enum):
@@ -304,3 +309,99 @@ def get_token_ids_logprobs_chunk(
pt += pruned_len
return next_split_pruned_len
def add_output_logprobs_for_spec_v1(
batch: ScheduleBatch,
res: Union[EagleVerifyOutput, NgramVerifyInput],
logits_output: Optional[LogitsProcessorOutput] = None,
):
# Extract args
if logits_output is None:
logits_output = res.logits_output
if hasattr(res, "accept_length_per_req_cpu"):
accept_length_per_req_cpu = res.accept_length_per_req_cpu
else:
# FIXME: Get a NgramVerifyOutput class and use that instead of this hack.
accept_length_per_req_cpu = res.accept_length.tolist()
top_logprobs_nums = batch.top_logprobs_nums
token_ids_logprobs = batch.token_ids_logprobs
accepted_indices = res.accepted_indices
assert len(accepted_indices) == len(logits_output.next_token_logits)
temperatures = batch.sampling_info.temperatures
num_draft_tokens = batch.spec_info.draft_token_num
# acceptance indices are the indices in a "flattened" batch.
# dividing it to num_draft_tokens will yield the actual batch index.
temperatures = temperatures[accepted_indices // num_draft_tokens]
if envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get():
logprobs = torch.nn.functional.log_softmax(
logits_output.next_token_logits, dim=-1
)
else:
logprobs = torch.nn.functional.log_softmax(
logits_output.next_token_logits / temperatures, dim=-1
)
batch_next_token_ids = res.verified_id
num_tokens_per_req = [accept + 1 for accept in accept_length_per_req_cpu]
# We should repeat top_logprobs_nums to match num_tokens_per_req.
top_logprobs_nums_repeat_interleaved = [
num
for num, num_tokens in zip(top_logprobs_nums, num_tokens_per_req)
for _ in range(num_tokens)
]
token_ids_logprobs_repeat_interleaved = [
token_ids
for token_ids, num_tokens in zip(token_ids_logprobs, num_tokens_per_req)
for _ in range(num_tokens)
]
# Extract logprobs
should_top_logprobs = any(x > 0 for x in top_logprobs_nums)
should_token_ids_logprobs = any(x is not None for x in token_ids_logprobs)
if should_top_logprobs:
(
logits_output.next_token_top_logprobs_val,
logits_output.next_token_top_logprobs_idx,
) = get_top_logprobs(
logprobs,
top_logprobs_nums_repeat_interleaved,
)
if should_token_ids_logprobs:
(
logits_output.next_token_token_ids_logprobs_val,
logits_output.next_token_token_ids_logprobs_idx,
) = get_token_ids_logprobs(
logprobs,
token_ids_logprobs_repeat_interleaved,
)
logits_output.next_token_logprobs = logprobs[
torch.arange(len(batch_next_token_ids), device=batch.sampling_info.device),
batch_next_token_ids,
]
# Add output logprobs to the request
pt = 0
next_token_logprobs = logits_output.next_token_logprobs.tolist()
verified_ids = batch_next_token_ids.tolist()
token_top_logprobs_val = logits_output.next_token_top_logprobs_val
token_top_logprobs_idx = logits_output.next_token_top_logprobs_idx
for req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True):
for _ in range(num_tokens):
# TODO: add token_ids_logprobs to each request
if req.return_logprob:
req.output_token_logprobs_val.append(next_token_logprobs[pt])
req.output_token_logprobs_idx.append(verified_ids[pt])
if req.top_logprobs_num > 0:
assert (
should_top_logprobs
), "Inconsistent state: should_top_logprobs is False"
req.output_top_logprobs_val.append(token_top_logprobs_val[pt])
req.output_top_logprobs_idx.append(token_top_logprobs_idx[pt])
pt += 1

85
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -14,7 +14,7 @@ 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 add_output_logprobs_for_spec_v1
from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
@@ -62,7 +62,6 @@ from sglang.srt.utils import (
MultiprocessingSerializer,
empty_context,
get_available_gpu_memory,
get_bool_env_var,
is_cuda,
is_npu,
next_power_of_2,
@@ -75,7 +74,6 @@ if is_cuda():
from sgl_kernel import segment_packbits # noqa: F401
logger = logging.getLogger(__name__)
SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB")
class EAGLEWorker(TpModelWorker):
@@ -756,7 +754,7 @@ class EAGLEWorker(TpModelWorker):
)
if batch.return_logprob:
self.add_logprob_values(batch, res, logits_output)
add_output_logprobs_for_spec_v1(batch, res, logits_output)
# Prepare the batch for the next draft forwards.
batch.forward_mode = (
@@ -845,85 +843,6 @@ class EAGLEWorker(TpModelWorker):
model=self.target_worker.model_runner.model,
)
def add_logprob_values(
self,
batch: ScheduleBatch,
res: EagleVerifyOutput,
logits_output: LogitsProcessorOutput,
):
# Extract args
logits_output = res.logits_output
top_logprobs_nums = batch.top_logprobs_nums
token_ids_logprobs = batch.token_ids_logprobs
accepted_indices = res.accepted_indices
assert len(accepted_indices) == len(logits_output.next_token_logits)
temperatures = batch.sampling_info.temperatures
num_draft_tokens = batch.spec_info.draft_token_num
# acceptance indices are the indices in a "flattened" batch.
# dividing it to num_draft_tokens will yield the actual batch index.
temperatures = temperatures[accepted_indices // num_draft_tokens]
if SGLANG_RETURN_ORIGINAL_LOGPROB:
logprobs = torch.nn.functional.log_softmax(
logits_output.next_token_logits, dim=-1
)
else:
logprobs = torch.nn.functional.log_softmax(
logits_output.next_token_logits / temperatures, dim=-1
)
batch_next_token_ids = res.verified_id
num_tokens_per_req = [accept + 1 for accept in res.accept_length_per_req_cpu]
# We should repeat top_logprobs_nums to match num_tokens_per_req.
top_logprobs_nums_repeat_interleaved = []
token_ids_logprobs_repeat_interleaved = []
for num, num_tokens in zip(top_logprobs_nums, num_tokens_per_req):
top_logprobs_nums_repeat_interleaved.extend([num] * num_tokens)
for token_ids, num_tokens in zip(token_ids_logprobs, num_tokens_per_req):
token_ids_logprobs_repeat_interleaved.extend([token_ids] * num_tokens)
# Extract logprobs
if any(x > 0 for x in top_logprobs_nums):
(
logits_output.next_token_top_logprobs_val,
logits_output.next_token_top_logprobs_idx,
) = get_top_logprobs(
logprobs,
top_logprobs_nums_repeat_interleaved,
)
if any(x is not None for x in token_ids_logprobs):
(
logits_output.next_token_token_ids_logprobs_val,
logits_output.next_token_token_ids_logprobs_idx,
) = get_token_ids_logprobs(
logprobs,
token_ids_logprobs_repeat_interleaved,
)
logits_output.next_token_logprobs = logprobs[
torch.arange(len(batch_next_token_ids), device=batch.sampling_info.device),
batch_next_token_ids,
]
# Add output logprobs to the request
pt = 0
next_token_logprobs = logits_output.next_token_logprobs.tolist()
verified_ids = batch_next_token_ids.tolist()
for req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True):
for _ in range(num_tokens):
if req.return_logprob:
req.output_token_logprobs_val.append(next_token_logprobs[pt])
req.output_token_logprobs_idx.append(verified_ids[pt])
if req.top_logprobs_num > 0:
req.output_top_logprobs_val.append(
res.logits_output.next_token_top_logprobs_val[pt]
)
req.output_top_logprobs_idx.append(
res.logits_output.next_token_top_logprobs_idx[pt]
)
pt += 1
def forward_draft_extend(
self,
batch: ScheduleBatch,

85
python/sglang/srt/speculative/multi_layer_eagle_worker.py
View File

@@ -22,7 +22,7 @@ from sglang.srt.distributed import get_tp_group
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_backend_context
from sglang.srt.layers.utils.logprob import get_token_ids_logprobs, get_top_logprobs
from sglang.srt.layers.utils.logprob import add_output_logprobs_for_spec_v1
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker
@@ -65,7 +65,6 @@ from sglang.srt.speculative.spec_utils import (
from sglang.srt.utils import (
empty_context,
get_available_gpu_memory,
get_bool_env_var,
is_cuda,
is_npu,
next_power_of_2,
@@ -77,7 +76,6 @@ if is_cuda():
from sgl_kernel import segment_packbits # noqa: F401
logger = logging.getLogger(__name__)
SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB")
class MultiLayerEagleWorker(TpModelWorker):
@@ -727,7 +725,7 @@ class MultiLayerEagleWorker(TpModelWorker):
)
if batch.return_logprob:
self.add_logprob_values(batch, res, logits_output)
add_output_logprobs_for_spec_v1(batch, res, logits_output)
# Prepare the batch for the next draft forwards.
batch.forward_mode = (
@@ -737,85 +735,6 @@ class MultiLayerEagleWorker(TpModelWorker):
return logits_output, res, model_worker_batch, can_run_cuda_graph
def add_logprob_values(
self,
batch: ScheduleBatch,
res: EagleVerifyOutput,
logits_output: LogitsProcessorOutput,
):
# Extract args
logits_output = res.logits_output
top_logprobs_nums = batch.top_logprobs_nums
token_ids_logprobs = batch.token_ids_logprobs
accepted_indices = res.accepted_indices
assert len(accepted_indices) == len(logits_output.next_token_logits)
temperatures = batch.sampling_info.temperatures
num_draft_tokens = batch.spec_info.draft_token_num
# acceptance indices are the indices in a "flattened" batch.
# dividing it to num_draft_tokens will yield the actual batch index.
temperatures = temperatures[accepted_indices // num_draft_tokens]
if SGLANG_RETURN_ORIGINAL_LOGPROB:
logprobs = torch.nn.functional.log_softmax(
logits_output.next_token_logits, dim=-1
)
else:
logprobs = torch.nn.functional.log_softmax(
logits_output.next_token_logits / temperatures, dim=-1
)
batch_next_token_ids = res.verified_id
num_tokens_per_req = [accept + 1 for accept in res.accept_length_per_req_cpu]
# We should repeat top_logprobs_nums to match num_tokens_per_req.
top_logprobs_nums_repeat_interleaved = []
token_ids_logprobs_repeat_interleaved = []
for num, num_tokens in zip(top_logprobs_nums, num_tokens_per_req):
top_logprobs_nums_repeat_interleaved.extend([num] * num_tokens)
for token_ids, num_tokens in zip(token_ids_logprobs, num_tokens_per_req):
token_ids_logprobs_repeat_interleaved.extend([token_ids] * num_tokens)
# Extract logprobs
if any(x > 0 for x in top_logprobs_nums):
(
logits_output.next_token_top_logprobs_val,
logits_output.next_token_top_logprobs_idx,
) = get_top_logprobs(
logprobs,
top_logprobs_nums_repeat_interleaved,
)
if any(x is not None for x in token_ids_logprobs):
(
logits_output.next_token_token_ids_logprobs_val,
logits_output.next_token_token_ids_logprobs_idx,
) = get_token_ids_logprobs(
logprobs,
token_ids_logprobs_repeat_interleaved,
)
logits_output.next_token_logprobs = logprobs[
torch.arange(len(batch_next_token_ids), device=batch.sampling_info.device),
batch_next_token_ids,
]
# Add output logprobs to the request
pt = 0
next_token_logprobs = logits_output.next_token_logprobs.tolist()
verified_ids = batch_next_token_ids.tolist()
for req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True):
for _ in range(num_tokens):
if req.return_logprob:
req.output_token_logprobs_val.append(next_token_logprobs[pt])
req.output_token_logprobs_idx.append(verified_ids[pt])
if req.top_logprobs_num > 0:
req.output_top_logprobs_val.append(
res.logits_output.next_token_top_logprobs_val[pt]
)
req.output_top_logprobs_idx.append(
res.logits_output.next_token_top_logprobs_idx[pt]
)
pt += 1
def forward_draft_extend(
self,
batch: ScheduleBatch,

32
python/sglang/srt/speculative/ngram_info.py
View File

@@ -155,7 +155,7 @@ class NgramVerifyInput(SpecInput):
batch: ScheduleBatch,
logits_output: torch.Tensor,
):
accept_index_cpu = self.accept_index.tolist()
accept_index_cpu = self.accepted_indices.tolist()
predict_cpu = self.predict.tolist()
has_finished = False
@@ -171,7 +171,7 @@ class NgramVerifyInput(SpecInput):
if req.finished():
has_finished = True
# set all tokens after finished token to -1 and break
self.accept_index[i, j + 1 :] = -1
self.accepted_indices[i, j + 1 :] = -1
break
else:
if req.grammar is not None:
@@ -180,7 +180,7 @@ class NgramVerifyInput(SpecInput):
except ValueError as e:
logger.info(
f"{i=}, {req=}\n"
f"{self.accept_index=}\n"
f"{self.accepted_indices=}\n"
f"{self.predict=}\n"
)
raise e
@@ -190,15 +190,17 @@ class NgramVerifyInput(SpecInput):
)
if has_finished:
self.accept_length = (self.accept_index != -1).sum(dim=1) - 1
self.accept_index = self.accept_index[self.accept_index != -1]
self.accept_length = (self.accepted_indices != -1).sum(dim=1) - 1
self.accepted_indices = self.accepted_indices[self.accepted_indices != -1]
logits_output.next_token_logits = logits_output.next_token_logits[
self.accept_index
self.accepted_indices
]
if logits_output.hidden_states:
logits_output.hidden_states = logits_output.hidden_states[self.accept_index]
self.verified_id = self.predict[self.accept_index]
logits_output.hidden_states = logits_output.hidden_states[
self.accepted_indices
]
self.verified_id = self.predict[self.accepted_indices]
def _free_cache(
self, batch: ScheduleBatch, page_size: int, accept_length_cpu: torch.Tensor
@@ -208,16 +210,16 @@ class NgramVerifyInput(SpecInput):
if page_size == 1:
# TODO: boolean array index leads to a device sync. Remove it.
evict_mask = torch.full_like(self.draft_token, True, dtype=torch.bool)
evict_mask[self.accept_index] = False
evict_mask[self.accepted_indices] = False
batch.token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask])
batch.out_cache_loc = batch.out_cache_loc[self.accept_index]
batch.out_cache_loc = batch.out_cache_loc[self.accepted_indices]
else:
# Shift the accepted tokens to the beginning.
# Only evict the last part
src_cache_loc, tgt_cache_loc, to_free_num_slots = get_src_tgt_cache_loc(
batch.seq_lens,
batch.out_cache_loc,
self.accept_index,
self.accepted_indices,
self.accept_length,
self.draft_token_num,
page_size,
@@ -285,14 +287,14 @@ class NgramVerifyInput(SpecInput):
predict_shape = list(logits_output.next_token_logits.shape)[:-1]
predict_shape[-1] += 1
self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device)
self.accept_index = torch.full(
self.accepted_indices = torch.full(
(bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device
)
self.accept_length = torch.empty((bs,), dtype=torch.int32, device=self.device)
verify_tree_greedy(
predicts=self.predict, # mutable
accept_index=self.accept_index, # mutable
accept_index=self.accepted_indices, # mutable
accept_token_num=self.accept_length, # mutable
candidates=candidates,
retrive_index=self.retrive_index,
@@ -312,7 +314,7 @@ class NgramVerifyInput(SpecInput):
predict_shape = list(logits_output.next_token_logits.shape)[:-1]
predict_shape[-1] += 1
self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device)
self.accept_index = torch.full(
self.accepted_indices = torch.full(
(bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device
)
self.accept_length = torch.empty((bs,), dtype=torch.int32, device=self.device)
@@ -354,7 +356,7 @@ class NgramVerifyInput(SpecInput):
)
tree_speculative_sampling_target_only(
predicts=self.predict, # mutable
accept_index=self.accept_index, # mutable
accept_index=self.accepted_indices, # mutable
accept_token_num=self.accept_length, # mutable
candidates=candidates.to(torch.int64),
retrive_index=self.retrive_index.to(torch.int64),

87
python/sglang/srt/speculative/ngram_worker.py
View File

@@ -5,9 +5,7 @@ import numpy as np
import torch
from sgl_kernel.speculative import reconstruct_indices_from_tree_mask
from sglang.srt.environ import envs
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.utils.logprob import get_token_ids_logprobs, get_top_logprobs
from sglang.srt.layers.utils.logprob import add_output_logprobs_for_spec_v1
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker
@@ -198,85 +196,6 @@ class NGRAMWorker:
)
batch.spec_info.prepare_for_verify(batch, self.page_size)
def add_logprob_values(
self,
batch: ScheduleBatch,
res: NgramVerifyInput,
logits_output: LogitsProcessorOutput,
):
# Extract args
top_logprobs_nums = batch.top_logprobs_nums
token_ids_logprobs = batch.token_ids_logprobs
accepted_indices = res.accept_index
assert len(accepted_indices) == len(logits_output.next_token_logits)
temperatures = batch.sampling_info.temperatures
num_draft_tokens = batch.spec_info.draft_token_num
# acceptance indices are the indices in a "flattened" batch.
# dividing it to num_draft_tokens will yield the actual batch index.
temperatures = temperatures[accepted_indices // num_draft_tokens]
if envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get():
logprobs = torch.nn.functional.log_softmax(
logits_output.next_token_logits, dim=-1
)
else:
logprobs = torch.nn.functional.log_softmax(
logits_output.next_token_logits / temperatures, dim=-1
)
batch_next_token_ids = res.verified_id
accept_length_per_req_cpu = res.accept_length.tolist()
num_tokens_per_req = [accept + 1 for accept in accept_length_per_req_cpu]
# We should repeat top_logprobs_nums to match num_tokens_per_req.
top_logprobs_nums_repeat_interleaved = []
token_ids_logprobs_repeat_interleaved = []
for num, num_tokens in zip(top_logprobs_nums, num_tokens_per_req):
top_logprobs_nums_repeat_interleaved.extend([num] * num_tokens)
for token_ids, num_tokens in zip(token_ids_logprobs, num_tokens_per_req):
token_ids_logprobs_repeat_interleaved.extend([token_ids] * num_tokens)
# Extract logprobs
if any(x > 0 for x in top_logprobs_nums):
(
logits_output.next_token_top_logprobs_val,
logits_output.next_token_top_logprobs_idx,
) = get_top_logprobs(
logprobs,
top_logprobs_nums_repeat_interleaved,
)
if any(x is not None for x in token_ids_logprobs):
(
logits_output.next_token_token_ids_logprobs_val,
logits_output.next_token_token_ids_logprobs_idx,
) = get_token_ids_logprobs(
logprobs,
token_ids_logprobs_repeat_interleaved,
)
logits_output.next_token_logprobs = logprobs[
torch.arange(len(batch_next_token_ids), device=batch.sampling_info.device),
batch_next_token_ids,
]
# Add output logprobs to the request
pt = 0
next_token_logprobs = logits_output.next_token_logprobs.tolist()
verified_ids = batch_next_token_ids.tolist()
for req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True):
for _ in range(num_tokens):
if req.return_logprob:
req.output_token_logprobs_val.append(next_token_logprobs[pt])
req.output_token_logprobs_idx.append(verified_ids[pt])
if req.top_logprobs_num > 0:
req.output_top_logprobs_val.append(
logits_output.next_token_top_logprobs_val[pt]
)
req.output_top_logprobs_idx.append(
logits_output.next_token_top_logprobs_idx[pt]
)
pt += 1
def _update_ngram_cache(self, batch: ScheduleBatch):
batch_tokens = []
for req in batch.reqs:
@@ -305,14 +224,14 @@ class NGRAMWorker:
batch_result.logits_output,
batch_result.can_run_cuda_graph,
)
verify_input = model_worker_batch.spec_info
verify_input: NgramVerifyInput = model_worker_batch.spec_info
logits_output, next_token_ids, num_accepted_tokens = verify_input.verify(
batch, logits_output, self.page_size
)
# Store accept_lens for per-request metrics
accept_lens = verify_input.accept_length
if batch.return_logprob:
self.add_logprob_values(batch, verify_input, logits_output)
add_output_logprobs_for_spec_v1(batch, verify_input, logits_output)
self._update_ngram_cache(batch)
batch.forward_mode = ForwardMode.DECODE