Fix the _chunked_req_scheduled_last_iter flag with a content-based stash gate (#26938)
This commit is contained in:
@@ -936,15 +936,6 @@ class Scheduler(
|
||||
elif self.chunked_prefill_size is not None and self.chunked_prefill_size <= 0:
|
||||
self.chunked_prefill_size = None
|
||||
self.chunked_req = None
|
||||
# Tracks whether the current self.chunked_req was actually scheduled
|
||||
# into last iteration's batch (i.e., in can_run_list -> got a fresh
|
||||
# req_pool_idx from prepare_for_extend). Used to gate the
|
||||
# stash_chunked_request call at the top of get_next_batch_to_run:
|
||||
# if add_chunked_req early-returned under hybrid-SWA pressure,
|
||||
# the req_pool_idx was already freed and the full_untruncated_fill_ids
|
||||
# was rebuilt by init_next_round_input, so running stash would
|
||||
# double-free and corrupt prefix_indices.
|
||||
self._chunked_req_scheduled_last_iter = False
|
||||
self.is_mixed_chunk = (
|
||||
self.chunked_prefill_size is not None
|
||||
and self.server_args.enable_mixed_chunk
|
||||
@@ -2443,7 +2434,11 @@ class Scheduler(
|
||||
# only finished requests to running_batch.
|
||||
chunked_req_to_exclude.add(self.chunked_req)
|
||||
|
||||
if self._chunked_req_scheduled_last_iter:
|
||||
# Stash (cache) the previous chunk only when it produced new KV
|
||||
# beyond what is already cached. A parked chunk (add_chunked_req
|
||||
# hybrid-SWA early-return) leaves fill_len == len(prefix_indices),
|
||||
# so there is nothing new to cache and stashing would be a no-op.
|
||||
if self.chunked_req.fill_len > len(self.chunked_req.prefix_indices):
|
||||
self.stash_chunked_request(self.chunked_req)
|
||||
|
||||
# HiSparse has its own prefill-to-decode transition; skip last_batch merge.
|
||||
@@ -2645,11 +2640,6 @@ class Scheduler(
|
||||
if self.chunked_req is not None:
|
||||
self.chunked_req.init_next_round_input()
|
||||
self.chunked_req = adder.add_chunked_req(self.chunked_req)
|
||||
self._chunked_req_scheduled_last_iter = (
|
||||
self.chunked_req in adder.can_run_list
|
||||
)
|
||||
else:
|
||||
self._chunked_req_scheduled_last_iter = False
|
||||
|
||||
if self.enable_lora:
|
||||
running_loras = {req.lora_id for req in self.running_batch.reqs}
|
||||
@@ -2749,9 +2739,6 @@ class Scheduler(
|
||||
# Update chunked prefill
|
||||
assert self.chunked_req is None
|
||||
self.chunked_req = adder.new_chunked_req
|
||||
# new_chunked_req is added to can_run_list by add_one_req,
|
||||
# so it will be scheduled this iter -> stash is needed next iter.
|
||||
self._chunked_req_scheduled_last_iter = True
|
||||
|
||||
if self.chunked_req is not None:
|
||||
self.chunked_req.inflight_middle_chunks += 1
|
||||
|
||||
Reference in New Issue
Block a user