From ed52d01b0b232f1b7c10e4b25fcc0f0fb7cae6ff Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 14 Dec 2025 00:32:54 +0800 Subject: [PATCH] Fix spec info's filter when reqs are finished right after prefill (#14742) --- python/sglang/srt/environ.py | 3 +++ python/sglang/srt/managers/schedule_batch.py | 11 +++++++---- python/sglang/srt/managers/scheduler.py | 6 +++--- python/sglang/srt/speculative/eagle_info.py | 11 ++++++++--- python/sglang/srt/speculative/eagle_worker.py | 15 --------------- test/registered/spec/eagle/test_eagle_infer_a.py | 16 ++++++++++++++++ 6 files changed, 37 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index fdf1fc901..ff380e002 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -321,6 +321,9 @@ class Envs: SGLANG_ENABLE_SPEC_V2 = EnvBool(False) SGLANG_ENABLE_OVERLAP_PLAN_STREAM = EnvBool(False) + # Spec Config + SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True) + # VLM SGLANG_VLM_CACHE_SIZE_MB = EnvInt(100) SGLANG_IMAGE_MAX_PIXELS = EnvInt(16384 * 28 * 28) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 3e788e6fb..5d0c44437 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1796,6 +1796,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self, chunked_req_to_exclude: Optional[Union[Req, List[Req]]] = None, keep_indices: Optional[List[int]] = None, + # FIXME(lsyin): deprecate this API after spec v1 is deprecated + v1_spec_info_filtered: Optional[bool] = False, ): # FIXME(lsyin): used here to get the correct seq_lens # The batch has been launched but we need it verified to get correct next batch info @@ -1852,11 +1854,12 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.has_grammar = any(req.grammar for req in self.reqs) self.sampling_info.filter_batch(keep_indices, keep_indices_device) + # NOTE: spec_info filtered before batch filtering only happens in: + # - Spec v1's verify phase + # - Only for decode batch (running_batch) + has_been_filtered = v1_spec_info_filtered and not self.is_v2_eagle + if self.spec_info: - if chunked_req_to_exclude is not None and len(chunked_req_to_exclude) > 0: - has_been_filtered = False - else: - has_been_filtered = True self.spec_info.filter_batch( new_indices=keep_indices_device, has_been_filtered=has_been_filtered, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 075f26df3..fdb817517 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1937,7 +1937,7 @@ class Scheduler( and not (new_batch.return_logprob or self.running_batch.return_logprob) ): # TODO (lianmin): support return_logprob + mixed chunked prefill - self.running_batch.filter_batch() + self.running_batch.filter_batch(v1_spec_info_filtered=True) if not self.running_batch.is_empty(): self.running_batch.prepare_for_decode() new_batch.mix_with_running(self.running_batch) @@ -1954,7 +1954,7 @@ class Scheduler( """Update the current running decoding batch.""" initial_bs = batch.batch_size() - batch.filter_batch() + batch.filter_batch(v1_spec_info_filtered=True) if batch.is_empty(): batch.batch_is_full = False return batch @@ -2509,7 +2509,7 @@ class Scheduler( self.cur_batch = None if recv_req.mode == "retract": - self.running_batch.filter_batch() + self.running_batch.filter_batch(v1_spec_info_filtered=True) if len(self.running_batch.reqs) != 0: retracted_reqs = self.running_batch.retract_all(self.server_args) for req in retracted_reqs: diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 2ffcda775..6566e0ce9 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -7,6 +7,7 @@ import torch import torch.nn.functional as F from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject +from sglang.srt.environ import envs from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.sampler import apply_custom_logit_processor @@ -754,13 +755,17 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): self.future_indices.indices = self.future_indices.indices[new_indices] return + strict_check = envs.SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK.get() if has_been_filtered: # in eagle_utils.py:verify, we have already filtered the batch by `unfinished_index` # therefore, we don't need to filter the batch again in scheduler + error_msg = f"length of new_indices: {len(new_indices)} != length of topk_p: {len(self.topk_p)}, this should not happen" if len(new_indices) != len(self.topk_p): - logger.warning( - f"length of new_indices: {len(new_indices)} != length of topk_p: {len(self.topk_p)}, this should not happen" - ) + if strict_check: + raise ValueError(error_msg) + else: + logger.warning(error_msg) + self.topk_p = self.topk_p[: len(new_indices)] self.topk_index = self.topk_index[: len(new_indices)] self.hidden_states = self.hidden_states[: len(new_indices)] diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 07e3798f1..8109640e8 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -909,21 +909,6 @@ class EAGLEWorker(TpModelWorker): assert isinstance(forward_batch.spec_info, EagleDraftInput) assert forward_batch.spec_info is batch.spec_info self.capture_for_decode(logits_output, forward_batch.spec_info) - has_finished, unfinished_req_index = False, [] - for i, req in enumerate(batch.reqs): - if req.finished(): - has_finished = True - else: - unfinished_req_index.append(i) - if has_finished: - unfinished_index_device = torch.tensor( - unfinished_req_index, - dtype=torch.int64, - device=batch.spec_info.topk_p.device, - ) - batch.spec_info.filter_batch( - unfinished_index_device, has_been_filtered=False - ) def forward_draft_extend_after_decode(self, batch: ScheduleBatch): assert isinstance(batch.spec_info, EagleDraftInput) diff --git a/test/registered/spec/eagle/test_eagle_infer_a.py b/test/registered/spec/eagle/test_eagle_infer_a.py index 0abaa1205..5c7e0d956 100644 --- a/test/registered/spec/eagle/test_eagle_infer_a.py +++ b/test/registered/spec/eagle/test_eagle_infer_a.py @@ -1,4 +1,5 @@ import os +import random import unittest import requests @@ -71,6 +72,7 @@ class TestEAGLEEngine(CustomTestCase): engine = sgl.Engine(**config, log_level="info", decode_log_interval=10) try: self._test_single_generation(engine) + self._test_first_token_finish(engine) self._test_batch_generation(engine) self._test_eos_token(engine) self._test_acc_length(engine) @@ -109,6 +111,20 @@ class TestEAGLEEngine(CustomTestCase): avg_spec_accept_length, self.THRESHOLDS["batch_avg_accept_len"] ) + def _test_first_token_finish(self, engine): + prompt = [ + f"There are {i} apples on the table. How to divide them equally?" + for i in range(8) + ] + params = [ + {"temperature": 0, "max_new_tokens": random.randint(1, 3)} for _ in range(8) + ] + outputs = engine.generate(prompt, params) + for i, output in enumerate(outputs): + print(f"Prompt: {prompt[i]}") + print(f"Generated: {output['text']}") + print("-" * 40) + def _test_eos_token(self, engine): prompt = "[INST] <>\nYou are a helpful assistant.\n<>\nToday is a sunny day and I like [/INST]" params = {