From 6350042696aa75a7ae1b4707a908ecbafa40015d Mon Sep 17 00:00:00 2001 From: Yixin Dong Date: Thu, 27 Nov 2025 04:31:46 -0800 Subject: [PATCH] feat: Naive support Spec V2 + Constrained Decoding (#13425) Signed-off-by: Ubospica Co-authored-by: Liangsheng Yin --- python/sglang/srt/managers/schedule_batch.py | 7 ++ python/sglang/srt/managers/scheduler.py | 13 +++- .../scheduler_output_processor_mixin.py | 10 ++- .../sglang/srt/speculative/eagle_info_v2.py | 8 ++ .../sglang/srt/speculative/eagle_worker_v2.py | 33 +++++++- .../sglang/test/kits/json_constrained_kit.py | 8 +- test/srt/run_suite.py | 1 + test/srt/test_eagle_constrained_decoding.py | 77 +++++++++++++++++++ 8 files changed, 149 insertions(+), 8 deletions(-) create mode 100644 test/srt/test_eagle_constrained_decoding.py diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 2637d57c4..a11e8fc7e 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1924,6 +1924,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): dimensions=self.dimensions, dllm_block_offsets=[req.dllm_block_offset for req in self.reqs], dllm_config=self.dllm_config, + reqs=self.reqs, + has_grammar=self.has_grammar, ) def copy(self): @@ -2041,3 +2043,8 @@ class ModelWorkerBatch: # Diffusion LLM dllm_block_offsets: Optional[List[int]] = None dllm_config: Optional[DllmConfig] = None + + # For constrained decoding + # FIXME(lsyin): remove this after fully overlap grammar + reqs: Optional[List[Req]] = None + has_grammar: bool = False diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 920e85292..61b774a03 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1016,7 +1016,16 @@ class Scheduler( and self.last_batch.forward_mode.is_extend() ) - if disable_overlap_for_batch: + # FIXME(lsyin): remove this grammar sync + need_grammar_sync = ( + batch is not None + and batch.forward_mode.is_decode() + and batch.has_grammar + and batch.is_v2_eagle + and len(self.result_queue) > 0 + ) + + if disable_overlap_for_batch or need_grammar_sync: pop_and_process() batch_result = None @@ -1025,7 +1034,7 @@ class Scheduler( self.result_queue.append((batch.copy(), batch_result)) if self.last_batch: - if not disable_overlap_for_batch: + if not disable_overlap_for_batch and not need_grammar_sync: pop_and_process() elif batch is None: # When the server is idle, do self-check and re-init some states diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 6997c09b3..c48f5f893 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -395,10 +395,16 @@ class SchedulerOutputProcessorMixin: logits_output.hidden_states[i].cpu().clone().tolist() ) - if req.grammar is not None and batch.spec_algorithm.is_none(): + if req.grammar is not None: # FIXME: this try-except block is for handling unexpected xgrammar issue. try: - req.grammar.accept_token(next_token_id) + if batch.spec_algorithm.is_none(): + # Normal decode: single token + req.grammar.accept_token(next_token_id) + elif batch.is_v2_eagle: + # Speculative decode: next_token_id is a list of accepted tokens + for token_id in next_token_id: + req.grammar.accept_token(token_id) except ValueError as e: # Grammar accept_token can raise ValueError if the token is not in the grammar. # This can happen if the grammar is not set correctly or the token is invalid. diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index fec3b00c3..952ce740c 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -256,6 +256,7 @@ class EagleVerifyInputV2Mixin: self: EagleVerifyInput, batch: ModelWorkerBatch, logits_output: LogitsProcessorOutput, + vocab_mask: torch.Tensor = None, ): """ Verify and find accepted tokens based on logits output and batch @@ -276,6 +277,13 @@ class EagleVerifyInputV2Mixin: next_token_logits = logits_output.next_token_logits device = batch.input_ids.device + # Apply grammar mask if provided + if vocab_mask is not None: + assert self.grammar is not None + self.grammar.apply_vocab_mask( + logits=next_token_logits, vocab_mask=vocab_mask + ) + candidates = self.draft_token.reshape(bs, self.draft_token_num) predict_shape = list(next_token_logits.shape)[:-1] predict = torch.zeros(predict_shape, dtype=torch.int32, device=device).flatten() diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 5ae77b990..665c551da 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -36,6 +36,7 @@ from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( detect_nan, draft_tp_context, + generate_token_bitmask, load_token_map, ) from sglang.srt.utils.common import ( @@ -667,7 +668,15 @@ class EAGLEWorkerV2(BaseSpecWorker): ), ) - # Run target verify batch in the main compute stream + # Prepare grammar data on CPU if needed + if batch.has_grammar: + retrieve_next_token_cpu = verify_input.retrive_next_token.cpu() + retrieve_next_sibling_cpu = verify_input.retrive_next_sibling.cpu() + draft_tokens_cpu = verify_input.draft_token.view( + verify_input.retrive_next_token.shape + ).cpu() + + # Run target verify batch in the main compute stream (GPU compute) forward_batch_output = self.target_worker.forward_batch_generation( model_worker_batch=None, forward_batch=verify_forward_batch, @@ -676,6 +685,26 @@ class EAGLEWorkerV2(BaseSpecWorker): ) logits_output = forward_batch_output.logits_output + # Generate vocab mask for constrained decoding + vocab_mask = None + if batch.has_grammar: + # Generate the logit mask for structured output. + vocab_mask = generate_token_bitmask( + batch.reqs, + verify_input, + retrieve_next_token_cpu, + retrieve_next_sibling_cpu, + draft_tokens_cpu, + batch.sampling_info.vocab_size, + ) + + if vocab_mask is not None: + assert verify_input.grammar is not None + vocab_mask = vocab_mask.to(verify_input.retrive_next_token.device) + # NOTE: otherwise, this vocab mask will be the one from the previous extend stage + # and will be applied to produce wrong results + batch.sampling_info.vocab_mask = None + # Sample if self.enable_nan_detection: detect_nan(logits_output) @@ -683,7 +712,7 @@ class EAGLEWorkerV2(BaseSpecWorker): predict, accept_length, accept_index, - ) = verify_input.sample(batch, logits_output) + ) = verify_input.sample(batch, logits_output, vocab_mask) new_seq_lens = batch.seq_lens + accept_length verify_done = torch.get_device_module(self.device).Event() verify_done.record() diff --git a/python/sglang/test/kits/json_constrained_kit.py b/python/sglang/test/kits/json_constrained_kit.py index 3e28e18f0..05c0b7f7d 100644 --- a/python/sglang/test/kits/json_constrained_kit.py +++ b/python/sglang/test/kits/json_constrained_kit.py @@ -24,7 +24,10 @@ class TestJSONConstrainedMixin: response = requests.post( self.base_url + "/generate", json={ - "text": "The capital of France is", + "text": ( + "Introduce the capital of France. Return in a JSON format. The JSON Schema is: " + + json.dumps(json_schema) + ), "sampling_params": { "temperature": 0 if n == 1 else 0.5, "max_new_tokens": 128, @@ -69,7 +72,8 @@ class TestJSONConstrainedMixin: {"role": "system", "content": "You are a helpful AI assistant"}, { "role": "user", - "content": "Introduce the capital of France. Return in a JSON format.", + "content": "Introduce the capital of France. Return in a JSON format. " + "The JSON Schema is: " + json.dumps(self.json_schema), }, ], temperature=0, diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 43295f394..1a0f89bf5 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -7,6 +7,7 @@ from sglang.test.ci.ci_utils import TestFile, run_unittest_files # NOTE: please sort the test cases alphabetically by the test file name suites = { "per-commit-1-gpu": [ + TestFile("test_eagle_constrained_decoding.py", 100), TestFile("debug_utils/test_tensor_dump_forward_hook.py", 15), TestFile("hicache/test_hicache_storage.py", 127), TestFile("hicache/test_hicache_variants.py", 393), diff --git a/test/srt/test_eagle_constrained_decoding.py b/test/srt/test_eagle_constrained_decoding.py new file mode 100644 index 000000000..f172618c7 --- /dev/null +++ b/test/srt/test_eagle_constrained_decoding.py @@ -0,0 +1,77 @@ +import unittest + +from sglang.srt.environ import envs +from sglang.srt.utils import kill_process_tree +from sglang.test.kits.json_constrained_kit import TestJSONConstrainedMixin +from sglang.test.kits.regex_constrained_kit import TestRegexConstrainedMixin +from sglang.test.test_utils import ( + DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST, + DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + + +class TestEagleConstrainedDecoding( + CustomTestCase, TestRegexConstrainedMixin, TestJSONConstrainedMixin +): + max_running_requests = 64 + attention_backend = "triton" + spec_steps = 5 + spec_topk = 1 + spec_draft_tokens = 6 + page_size = 1 + other_launch_args = [] + model = DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST + draft_model = DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST + grammar_backend = "xgrammar" + eagle_v2 = False + + @classmethod + def setUpClass(cls): + cls.base_url = DEFAULT_URL_FOR_TEST + launch_args = [ + "--trust-remote-code", + "--attention-backend", + cls.attention_backend, + "--speculative-algorithm", + "EAGLE", + "--speculative-draft-model", + cls.draft_model, + "--speculative-num-steps", + cls.spec_steps, + "--speculative-eagle-topk", + cls.spec_topk, + "--speculative-num-draft-tokens", + cls.spec_draft_tokens, + "--page-size", + str(cls.page_size), + "--mem-fraction-static", + "0.75", + "--max-running-requests", + str(cls.max_running_requests), + "--grammar-backend", + cls.grammar_backend, + ] + launch_args.extend(cls.other_launch_args) + with envs.SGLANG_ENABLE_SPEC_V2.override(cls.eagle_v2): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=launch_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + +class TestEagleConstrainedDecodingV2(TestEagleConstrainedDecoding): + eagle_v2 = True + + +if __name__ == "__main__": + unittest.main()