Fix Req fill_len (fill_ids) having dual semantics by restricting to truncated/committed semantics (#26659)

This commit is contained in:
fzyzcjy
2026-06-08 14:52:48 +08:00
committed by GitHub
parent 259a2da3e0
commit 9034c2f9ae
4 changed files with 21 additions and 11 deletions
+2 -2
View File
@@ -42,11 +42,11 @@ class ReqDllmMixin:
prefix_length = len(self.prefix_indices) prefix_length = len(self.prefix_indices)
min_required_length = prefix_length + self.dllm_config.block_size 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 # still incoming stage
return 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 is_prefill_phase = self.dllm_config.mask_id not in input_block
if is_prefill_phase: if is_prefill_phase:
+5 -5
View File
@@ -1078,9 +1078,8 @@ class Req(ReqDllmMixin):
self.determine_dllm_phase() self.determine_dllm_phase()
else: else:
self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids 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 # Streaming sessions reuse committed KV from the session slot, so
# custom logprob_start_len is not supported — override to -1. # custom logprob_start_len is not supported — override to -1.
@@ -1098,7 +1097,7 @@ class Req(ReqDllmMixin):
) )
self.logprob_start_len = -1 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) : 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: 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 # 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.swa_evicted_seqlen = 0
self.extend_batch_idx = 0 self.extend_batch_idx = 0
self.decode_batch_idx = 0 self.decode_batch_idx = 0
self.fill_len = 0
# When using input_embeds, we cannot easily mix the original input embeddings # 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. # 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 # - extend_input_len: Number of tokens that need to be processed in this extend batch
self.extend_input_len = extend_input_len self.extend_input_len = extend_input_len
if self.logprob_start_len == -1: if self.logprob_start_len == -1:
logprob_start_len = self.fill_len logprob_start_len = len(self.full_untruncated_fill_ids)
else: else:
# logprob_start_len should be at least the length of the prefix indices # 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)) logprob_start_len = max(self.logprob_start_len, len(self.prefix_indices))
+13 -3
View File
@@ -827,7 +827,11 @@ class PrefillAdder:
self.rem_chunk_tokens is None # chunked prefill is disabled self.rem_chunk_tokens is None # chunked prefill is disabled
or req.extend_input_len <= self.rem_chunk_tokens # it is the last chunk 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.can_run_list.append(req)
self._update_prefill_budget( self._update_prefill_budget(
0, 0,
@@ -932,7 +936,9 @@ class PrefillAdder:
) )
) )
req.prefix_indices = torch.cat([req.prefix_indices, new_indices]) 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) prefix_len = len(req.prefix_indices)
req.cache_protected_len = prefix_len req.cache_protected_len = prefix_len
@@ -959,7 +965,11 @@ class PrefillAdder:
self._add_dllm_req(req, prefix_len) self._add_dllm_req(req, prefix_len)
self._req_inc_lock_ref(req) self._req_inc_lock_ref(req)
elif self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens: 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.can_run_list.append(req)
self._req_inc_lock_ref(req) self._req_inc_lock_ref(req)
+1 -1
View File
@@ -2159,7 +2159,7 @@ class Scheduler(
if last_host_node.backuped or last_host_node is self.tree_cache.root_node: if last_host_node.backuped or last_host_node is self.tree_cache.root_node:
last_hash = last_host_node.get_last_hash_value() last_hash = last_host_node.get_last_hash_value()
matched_len = len(req.prefix_indices) + req.host_hit_length 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 = ( prefix_keys = (
last_host_node.get_prefix_hash_values(last_host_node.parent) last_host_node.get_prefix_hash_values(last_host_node.parent)