minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Isolate spec V1 path in decode post-processing (#22146)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-04-05 03:16:56 -07:00
committed by GitHub
parent 106baedbfb
commit cd2d45e220
5 changed files with 60 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/model_config.py
+2
View File
@@ -232,6 +232,8 @@ class ModelConfig:
# Cache attributes
self.hf_eos_token_id = self._get_hf_eos_token_id()
# Set by scheduler when reasoning_parser is enabled
self.think_end_id: Optional[int] = None
# multimodal
self.image_token_id = getattr(
python/sglang/srt/managers/scheduler.py
+1
View File
@@ -553,6 +553,7 @@ class Scheduler(
reasoning_parser.detector.think_end_token, add_special_tokens=False
)[0]
self._think_end_id = self.tokenizer.think_end_id
self.model_config.think_end_id = self._think_end_id
else:
self._think_end_id = None
python/sglang/srt/managers/scheduler_output_processor_mixin.py
+51 -50
View File
@@ -406,20 +406,8 @@ class SchedulerOutputProcessorMixin:
v.tolist()
for v in logits_output.next_token_token_ids_logprobs_val
]
else:
# for normal spec decoding: unify next_token_ids format
next_token_ids = []
cum_num_tokens = 0
next_token_ids_list = result.next_token_ids.tolist()
for i, req in enumerate(batch.reqs):
accept_length = result.accept_length_per_req_cpu[i]
next_token_ids.append(
next_token_ids_list[
cum_num_tokens : cum_num_tokens + accept_length + 1
]
)
cum_num_tokens += accept_length + 1
# else: Spec V1 — output_ids, check_finished, grammar, and reasoning tokens
# are already handled in the verify phase (eagle_info.py / ngram_info.py).
self.num_generated_tokens += len(batch.reqs)
if not batch.spec_algorithm.is_none():
@@ -431,24 +419,36 @@ class SchedulerOutputProcessorMixin:
self.token_to_kv_pool_allocator.free_group_begin()
# NOTE: in any case, we should check finish here
# if finished, also clean up committed kv cache and over-allocated kv cache here
# Spec V1 handles output_ids, check_finished, grammar, and reasoning tokens
# in the verify phase. Non-spec and V2 handle them here in post-processing.
is_spec_v1 = not batch.spec_algorithm.is_none() and not batch.is_spec_v2
# Check finish condition
for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
for i, req in enumerate(batch.reqs):
req: Req
if self.enable_overlap and (req.finished() or req.is_retracted):
# NOTE: This (req.finished() or req.is_retracted) should only happen when overlap scheduling is enabled.
# (currently not, e.g. Eagle V1 still check finish during forward)
# And all the over-allocated tokens will be freed in `release_kv_cache`.
continue
if is_spec_v1:
self._mamba_prefix_cache_update(req, batch, result, i)
req.time_stats.set_last_decode_finish_time()
self._handle_finished_req(req, i, logits_output)
if req.return_hidden_states and logits_output.hidden_states is not None:
req.hidden_states.append(
logits_output.hidden_states[i].cpu().clone().tolist()
)
if req.grammar is not None:
req.grammar.finished = req.finished()
continue
# Non-spec and V2: full post-processing
next_token_id = next_token_ids[i]
new_accepted_len = 1
if batch.spec_algorithm.is_none():
req.output_ids.append(next_token_id)
elif batch.is_spec_v2:
# Only spec v2's output_ids are updated here.
else:
req.output_ids.extend(next_token_id)
new_accepted_len = len(next_token_id)
@@ -456,39 +456,12 @@ class SchedulerOutputProcessorMixin:
# Update Mamba last track seqlen
self._mamba_prefix_cache_update(req, batch, result, i)
req.time_stats.set_last_decode_finish_time()
req.check_finished(new_accepted_len)
if (
self.server_args.disaggregation_decode_enable_offload_kvcache
and not req.finished()
):
self.decode_offload_manager.offload_kv_cache(req)
self._handle_finished_req(req, i, logits_output)
if req.finished():
# delete feature to save memory
if req.multimodal_inputs is not None and req.session is None:
req.multimodal_inputs.release_features()
self.maybe_collect_routed_experts(req)
if self.server_args.disaggregation_decode_enable_offload_kvcache:
# Asynchronously offload KV cache; release_kv_cache will be called after Device->Host transfer completes
if not self.decode_offload_manager.offload_kv_cache(req):
self.decode_offload_manager.finalize_release_on_finish(req)
else:
if self.enable_hisparse:
self.hisparse_coordinator.request_finished(req)
release_kv_cache(req, self.tree_cache)
req.time_stats.set_completion_time()
self.maybe_collect_customized_info(i, req, logits_output)
if req.return_logprob and (
batch.spec_algorithm.is_none() or batch.is_spec_v2
):
if req.return_logprob:
# Spec v1 handles logprobs inside its own worker.
# Normalize: non-spec has 1 token, spec v2 has multiple.
if batch.is_spec_v2:
@@ -554,6 +527,34 @@ class SchedulerOutputProcessorMixin:
num_accepted_tokens=result.num_accepted_tokens,
)
def _handle_finished_req(
self: Scheduler, req: Req, i: int, logits_output: LogitsProcessorOutput
):
if (
self.server_args.disaggregation_decode_enable_offload_kvcache
and not req.finished()
):
self.decode_offload_manager.offload_kv_cache(req)
if req.finished():
# delete feature to save memory
if req.multimodal_inputs is not None and req.session is None:
req.multimodal_inputs.release_features()
self.maybe_collect_routed_experts(req)
if self.server_args.disaggregation_decode_enable_offload_kvcache:
# Asynchronously offload KV cache; release_kv_cache will be called after Device->Host transfer completes
if not self.decode_offload_manager.offload_kv_cache(req):
self.decode_offload_manager.finalize_release_on_finish(req)
else:
if self.enable_hisparse:
self.hisparse_coordinator.request_finished(req)
release_kv_cache(req, self.tree_cache)
req.time_stats.set_completion_time()
self.maybe_collect_customized_info(i, req, logits_output)
def _maybe_update_reasoning_tokens(
self: Scheduler, req: Req, next_token_id: Union[int, List[int]]
):
python/sglang/srt/speculative/eagle_info.py
+3
View File
@@ -392,6 +392,7 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
accept_index_cpu = accept_index.tolist()
predict_cpu = predict.tolist()
has_finished = False
think_end_id = batch.model_config.think_end_id
# Iterate every accepted token and check if req has finished after append the token
# should be checked BEFORE free kv cache slots
@@ -403,6 +404,8 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
num_accepted += 1
id = predict_cpu[idx]
req.output_ids.append(id)
if req.require_reasoning and think_end_id is not None:
req.update_reasoning_tokens(id, think_end_id)
req.check_finished()
if req.finished():
has_finished = True
python/sglang/srt/speculative/ngram_info.py
+3
View File
@@ -161,6 +161,7 @@ class NgramVerifyInput(SpecInput):
accept_index_cpu = self.accepted_indices.tolist()
predict_cpu = self.predict.tolist()
has_finished = False
think_end_id = batch.model_config.think_end_id
# Iterate every accepted token and check if req has finished after append the token
# should be checked BEFORE free kv cache slots
@@ -170,6 +171,8 @@ class NgramVerifyInput(SpecInput):
break
id = predict_cpu[idx]
req.output_ids.append(id)
if req.require_reasoning and think_end_id is not None:
req.update_reasoning_tokens(id, think_end_id)
req.check_finished()
if req.finished():
has_finished = True