From 8fe8b6357601a768ab9bb0747bd802a21da4f8ba Mon Sep 17 00:00:00 2001 From: Stefan He Date: Mon, 1 Dec 2025 18:18:03 -0800 Subject: [PATCH] Revert "Try to remove wrong logic about max total token in spec decoding" (#14259) --- .../sglang/srt/model_executor/model_runner.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 6642038b1..c754bc54d 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1697,8 +1697,25 @@ class ModelRunner: if self.spec_algorithm.is_eagle() or self.spec_algorithm.is_standalone(): if self.is_draft_worker: + self.max_total_num_tokens = self.server_args.draft_runner_cache_size max_num_reqs = self.server_args.max_num_reqs else: + # We are sharing the `token_to_kv_pool`, and both verify and draft tokens + # can be concurrently allocated, so we should give a headroom for it. + self.server_args.draft_runner_cache_size = ( + self.max_total_num_tokens + # draft + + max_num_reqs + * self.server_args.speculative_num_steps + * self.server_args.speculative_eagle_topk + # verify + + max_num_reqs * self.server_args.speculative_num_draft_tokens + # buffer + + 100 + ) + # Target worker and draft worker shares the same indices for the + # token_to_kv_pool, so we should make sure to match max_total_num_tokens. + self.max_total_num_tokens = self.server_args.draft_runner_cache_size self.server_args.max_num_reqs = max_num_reqs if max_total_tokens is not None: