diff --git a/python/sglang/srt/dllm/mixin/req.py b/python/sglang/srt/dllm/mixin/req.py index 1c9cfc3b8..917893f56 100644 --- a/python/sglang/srt/dllm/mixin/req.py +++ b/python/sglang/srt/dllm/mixin/req.py @@ -42,11 +42,11 @@ class ReqDllmMixin: prefix_length = len(self.prefix_indices) min_required_length = prefix_length + self.dllm_config.block_size - if self.fill_len < min_required_length: + if len(self.full_untruncated_fill_ids) < min_required_length: # still incoming stage return - input_block = self.get_fill_ids()[prefix_length:min_required_length] + input_block = self.full_untruncated_fill_ids[prefix_length:min_required_length] is_prefill_phase = self.dllm_config.mask_id not in input_block if is_prefill_phase: diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 3c6ad40aa..c51a27f2f 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1078,9 +1078,8 @@ class Req(ReqDllmMixin): self.determine_dllm_phase() else: self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids - self.fill_len = len(self.full_untruncated_fill_ids) - input_len = self.fill_len + input_len = len(self.full_untruncated_fill_ids) # Streaming sessions reuse committed KV from the session slot, so # custom logprob_start_len is not supported — override to -1. @@ -1098,7 +1097,7 @@ class Req(ReqDllmMixin): ) self.logprob_start_len = -1 - token_ids_to_match = self.get_fill_ids()[ + token_ids_to_match = self.full_untruncated_fill_ids[ : self._compute_max_prefix_len(input_len) ] @@ -1163,7 +1162,7 @@ class Req(ReqDllmMixin): ) ) - self.set_extend_input_len(self.fill_len - len(self.prefix_indices)) + self.set_extend_input_len(input_len - len(self.prefix_indices)) def _compute_max_prefix_len(self, input_len: int) -> int: # NOTE: the matched length is at most 1 less than the input length to enable logprob computation @@ -1371,6 +1370,7 @@ class Req(ReqDllmMixin): self.swa_evicted_seqlen = 0 self.extend_batch_idx = 0 self.decode_batch_idx = 0 + self.fill_len = 0 # When using input_embeds, we cannot easily mix the original input embeddings # with the newly generated output token IDs during re-prefill of retracted request. @@ -1429,7 +1429,7 @@ class Req(ReqDllmMixin): # - extend_input_len: Number of tokens that need to be processed in this extend batch self.extend_input_len = extend_input_len if self.logprob_start_len == -1: - logprob_start_len = self.fill_len + logprob_start_len = len(self.full_untruncated_fill_ids) else: # logprob_start_len should be at least the length of the prefix indices logprob_start_len = max(self.logprob_start_len, len(self.prefix_indices)) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 9fae0bab4..9841686a1 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -827,7 +827,11 @@ class PrefillAdder: self.rem_chunk_tokens is None # chunked prefill is disabled or req.extend_input_len <= self.rem_chunk_tokens # it is the last chunk ): - # Non-chunked prefill + # Non-chunked prefill — the whole sequence is committed this iter. + req.fill_len = len(req.full_untruncated_fill_ids) + assert ( + req.fill_len == len(req.prefix_indices) + req.extend_input_len + ), f"{req.fill_len=} {len(req.prefix_indices)=} {req.extend_input_len=}" self.can_run_list.append(req) self._update_prefill_budget( 0, @@ -932,7 +936,9 @@ class PrefillAdder: ) ) req.prefix_indices = torch.cat([req.prefix_indices, new_indices]) - req.set_extend_input_len(req.fill_len - len(req.prefix_indices)) + req.set_extend_input_len( + len(req.full_untruncated_fill_ids) - len(req.prefix_indices) + ) prefix_len = len(req.prefix_indices) req.cache_protected_len = prefix_len @@ -959,7 +965,11 @@ class PrefillAdder: self._add_dllm_req(req, prefix_len) self._req_inc_lock_ref(req) elif self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens: - # Non-chunked prefill + # Non-chunked prefill — the whole sequence is committed this iter. + req.fill_len = len(req.full_untruncated_fill_ids) + assert ( + req.fill_len == len(req.prefix_indices) + req.extend_input_len + ), f"{req.fill_len=} {len(req.prefix_indices)=} {req.extend_input_len=}" self.can_run_list.append(req) self._req_inc_lock_ref(req) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index bbd32aaed..cb67a7755 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2159,7 +2159,7 @@ class Scheduler( if last_host_node.backuped or last_host_node is self.tree_cache.root_node: last_hash = last_host_node.get_last_hash_value() matched_len = len(req.prefix_indices) + req.host_hit_length - new_input_tokens = req.get_fill_ids()[matched_len:] + new_input_tokens = req.full_untruncated_fill_ids[matched_len:] prefix_keys = ( last_host_node.get_prefix_hash_values(last_host_node.parent)