[Fix] Lock PP dynamic-chunk profiling requests before releasing through the tree cache (#37509)

This commit is contained in:
Liangsheng Yin
2026-09-02 02:09:01 -07:00
committed by GitHub
parent ebfd8c60e5
commit 862a909a08
3 changed files with 8 additions and 6 deletions
+2 -3
View File
@@ -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).
+1 -1
View File
@@ -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)
)