From 6f08488042eae6f05dccc6d7d5082a9e4c5a07b1 Mon Sep 17 00:00:00 2001 From: Zhihao Zhang <44469583+a4zhangfei@users.noreply.github.com> Date: Mon, 10 Nov 2025 18:04:42 +0800 Subject: [PATCH] fix missing output_token_logprobs when using ngram speculative decoding (#10702) Co-authored-by: a4zhangfei --- .../srt/speculative/cpp_ngram/ngram.cpp | 7 ++ .../sglang/srt/speculative/cpp_ngram/ngram.h | 1 + .../sglang/srt/speculative/cpp_ngram/param.h | 1 + .../sglang/srt/speculative/cpp_ngram/queue.h | 2 + python/sglang/srt/speculative/ngram_worker.py | 86 +++++++++++++++++++ 5 files changed, 97 insertions(+) diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram.cpp b/python/sglang/srt/speculative/cpp_ngram/ngram.cpp index d1e982358..e7f0297e2 100644 --- a/python/sglang/srt/speculative/cpp_ngram/ngram.cpp +++ b/python/sglang/srt/speculative/cpp_ngram/ngram.cpp @@ -1,9 +1,16 @@ #include "ngram.h" #include +#include #include #include +#include +#include #include +#include +#include +#include +#include #include namespace ngram { diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram.h b/python/sglang/srt/speculative/cpp_ngram/ngram.h index bf0af0df9..3c9a9380e 100644 --- a/python/sglang/srt/speculative/cpp_ngram/ngram.h +++ b/python/sglang/srt/speculative/cpp_ngram/ngram.h @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include diff --git a/python/sglang/srt/speculative/cpp_ngram/param.h b/python/sglang/srt/speculative/cpp_ngram/param.h index 967832ad6..08b975bb1 100644 --- a/python/sglang/srt/speculative/cpp_ngram/param.h +++ b/python/sglang/srt/speculative/cpp_ngram/param.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include #include diff --git a/python/sglang/srt/speculative/cpp_ngram/queue.h b/python/sglang/srt/speculative/cpp_ngram/queue.h index e84a0fa7b..f08aa5cec 100644 --- a/python/sglang/srt/speculative/cpp_ngram/queue.h +++ b/python/sglang/srt/speculative/cpp_ngram/queue.h @@ -1,7 +1,9 @@ #pragma once #include +#include #include +#include namespace utils { diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index e1676ad1e..b6a1bf554 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -5,6 +5,8 @@ import numpy as np import torch from sgl_kernel.speculative import reconstruct_indices_from_tree_mask +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.sampler import get_token_ids_logprobs, get_top_logprobs from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker @@ -13,9 +15,12 @@ from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.cpp_ngram.ngram_cache import NgramCache from sglang.srt.speculative.ngram_info import NgramVerifyInput from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.utils import get_bool_env_var logger = logging.getLogger(__name__) +RETURN_ORIGINAL_LOGPROB = get_bool_env_var("RETURN_ORIGINAL_LOGPROB") + USE_FULL_MASK = True @@ -194,6 +199,85 @@ 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 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 + 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: @@ -225,6 +309,8 @@ class NGRAMWorker: logits_output, next_token_ids, num_accepted_tokens = verify_input.verify( batch, logits_output, self.page_size ) + if batch.return_logprob: + self.add_logprob_values(batch, verify_input, logits_output) self._update_ngram_cache(batch) batch.forward_mode = ForwardMode.DECODE