Fix Req fill_len (fill_ids) having dual semantics by restricting to truncated/committed semantics (#26659)
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user