[Fix] Lock PP dynamic-chunk profiling requests before releasing through the tree cache (#37509)
This commit is contained in:
@@ -1371,9 +1371,8 @@ class Req(ReqDllmMixin):
|
||||
appending only the new output tokens.
|
||||
|
||||
Falls back to a full rebuild when the in-place append is invalid:
|
||||
- aliasing: scheduler_pp_mixin assigns full_untruncated_fill_ids =
|
||||
origin_input_ids directly, so extending in place would write output
|
||||
tokens into the origin;
|
||||
- aliasing: full_untruncated_fill_ids is origin_input_ids itself, so
|
||||
extending in place would write output tokens into the origin;
|
||||
- lengths disagree: fresh req (array still empty), retraction
|
||||
(output_ids reset to empty), or set_finish_with_abort (origin
|
||||
replaced by a 1-token stub).
|
||||
|
||||
@@ -1253,7 +1253,7 @@ class Scheduler(
|
||||
self.profile_and_init_predictor()
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"[PP Dynamic Chunk] Failed to profile prefill latency: {e}. "
|
||||
f"[PP Dynamic Chunk] Failed to profile prefill latency: {e!r}. "
|
||||
"Dynamic chunking will be disabled."
|
||||
)
|
||||
self.enable_dynamic_chunking = False
|
||||
|
||||
@@ -634,8 +634,11 @@ class SchedulerPPMixin:
|
||||
origin_input_ids=input_ids,
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
req.full_untruncated_fill_ids = req.origin_input_ids
|
||||
req.logprob_start_len = -1
|
||||
# Walk the same match -> lock -> alloc lifecycle as a scheduled
|
||||
# request so release_kv_cache can release it symmetrically.
|
||||
req.init_next_round_input(self.tree_cache)
|
||||
lock = self.tree_cache.inc_lock_ref(req.last_node)
|
||||
req.swa_uuid_for_lock = lock.swa_uuid_for_lock
|
||||
req.set_extend_range(
|
||||
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user