From 424591d53d1e0becbd4df826d7995a81006604e7 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Tue, 30 Sep 2025 14:44:02 +0800 Subject: [PATCH] Fix spec filter batch when target extend (#10991) --- python/sglang/srt/managers/schedule_batch.py | 9 ++++++++- python/sglang/srt/speculative/ngram_utils.py | 2 +- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 01859b6ef..97f5d00ff 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1736,7 +1736,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.sampling_info.filter_batch(keep_indices, keep_indices_device) if self.spec_info: - self.spec_info.filter_batch(keep_indices_device) + 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, + ) def merge_batch(self, other: "ScheduleBatch"): # Penalizer orchestrator must be merged before Batch.reqs is merged. This is because diff --git a/python/sglang/srt/speculative/ngram_utils.py b/python/sglang/srt/speculative/ngram_utils.py index d675d35b5..d0e80c0a4 100644 --- a/python/sglang/srt/speculative/ngram_utils.py +++ b/python/sglang/srt/speculative/ngram_utils.py @@ -405,7 +405,7 @@ class NgramVerifyInput: return logits_output, self.verified_id, self.accept_length.sum().item() - def filter_batch(self, new_indices: torch.Tensor): + def filter_batch(self, new_indices: torch.Tensor, has_been_filtered: bool = True): pass def merge_batch(self, spec_info: NgramVerifyInput):