diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index aff05bf42..7bd3dc3e6 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -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 diff --git a/python/sglang/srt/layers/utils/logprob.py b/python/sglang/srt/layers/utils/logprob.py index 6f84c15ff..5504008a0 100644 --- a/python/sglang/srt/layers/utils/logprob.py +++ b/python/sglang/srt/layers/utils/logprob.py @@ -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, ) diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 53c4e31c0..f4b2b9978 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -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( diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 9eea302b7..ba7773300 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index c32d8a83f..f3418db13 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -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, diff --git a/test/registered/spec/eagle/test_eagle_infer_beta.py b/test/registered/spec/eagle/test_eagle_infer_beta.py index 313fd5512..3124877a9 100644 --- a/test/registered/spec/eagle/test_eagle_infer_beta.py +++ b/test/registered/spec/eagle/test_eagle_infer_beta.py @@ -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"]