fix missing output_token_logprobs when using ngram speculative decoding (#10702)

Co-authored-by: a4zhangfei <a4zhangfei@qq.com>

python/sglang/srt/speculative/cpp_ngram/ngram.cpp

@@ -1,9 +1,16 @@
 #include "ngram.h"
 
 #include <algorithm>
+#include <chrono>
 #include <cstring>
 #include <limits>
+#include <list>
+#include <mutex>
 #include <queue>
+#include <stdexcept>
+#include <thread>
+#include <tuple>
+#include <unordered_map>
 #include <vector>
 
 namespace ngram {

python/sglang/srt/speculative/cpp_ngram/ngram.h

@@ -5,6 +5,7 @@
 #include <functional>
 #include <list>
 #include <mutex>
+#include <new>
 #include <set>
 #include <sstream>
 #include <thread>

python/sglang/srt/speculative/cpp_ngram/param.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <cstddef>
+#include <cstdlib>
 #include <iostream>
 #include <limits>
 #include <regex>

python/sglang/srt/speculative/cpp_ngram/queue.h

@@ -1,7 +1,9 @@
 #pragma once
 
 #include <condition_variable>
+#include <mutex>
 #include <queue>
+#include <utility>
 
 namespace utils {
 

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