From e864ded5d8c485814a3c3a590ac90660e2a0780c Mon Sep 17 00:00:00 2001 From: LyxWxj Date: Thu, 29 Oct 2026 00:27:30 +0800 Subject: [PATCH 1/6] refactor(scheduler): unify prefill progress and wavefront tasks Share wavefront prefill progress handling between output and submission updates. Unify PRELUDE and RECURRENT wavefront task admission while preserving candidate queue order, token budgeting, capacity checks, and stable task ordering. Signed-off-by: LyxWxj --- vllm_rlt/core/scheduler.py | 241 +++++++++++++++++++------------------ 1 file changed, 122 insertions(+), 119 deletions(-) diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index 62e0cdee..9d5cd804 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -882,63 +882,55 @@ def _take_wavefront_phase( added = set() blocked = set() - if phase == RequestPhase.PREFILLING_PRELUDE: - cursors = {request.request_id: request.num_preluded_tokens for request in candidates} - for _ in range(len(candidates[0].prefill_depth_ranges)): - for request in candidates: - request_id = request.request_id - if token_budget <= 0: - break - if request_id in blocked or request.phase != phase: - continue - start = cursors[request_id] - if start >= len(request.prompt_token_ids): - continue - scheduled_request = self._make_wavefront_task( - request, phase, token_budget, start=start - ) - if scheduled_request is None: - continue - frontier = scheduled_request.token_start + scheduled_request.token_count - if not self._ensure_execution_capacity(request, frontier, selected=selected): - blocked.add(request_id) - continue - scheduled_requests.append(replace(scheduled_request, phase=phase)) - selected.add(request_id) - added.add(request_id) - cursors[request_id] += scheduled_request.token_count - token_budget -= scheduled_request.token_count - else: - tasks = [] - for depth in range(len(candidates[0].prefill_depth_ranges)): - for owner_index, request in enumerate(candidates): - request_id = request.request_id - if token_budget <= 0: - break - if request_id in blocked or request.phase != phase: - continue + total_depths = len(candidates[0].prefill_depth_ranges) + + def accept_task(request: Request, scheduled_request: ScheduledRequest) -> bool: + """Reserve capacity and commit one wavefront task to this batch.""" + nonlocal token_budget + request_id = request.request_id + frontier = scheduled_request.token_start + scheduled_request.token_count + if not self._ensure_execution_capacity(request, frontier, selected=selected): + blocked.add(request_id) + return False + selected.add(request_id) + added.add(request_id) + token_budget -= scheduled_request.token_count + return True + + tasks = [] + prelude_cursors = { + request.request_id: request.num_preluded_tokens for request in candidates + } + for depth in range(total_depths): + for owner_index, request in enumerate(candidates): + request_id = request.request_id + if token_budget <= 0: + break + if request_id in blocked or request.phase != phase: + continue + if phase == RequestPhase.PREFILLING_PRELUDE: + start, end = prelude_cursors[request_id], len(request.prompt_token_ids) + else: start, end = request.prefill_depth_ranges[depth] end = min(end, request.num_preluded_tokens) - if start >= end: - continue - scheduled_request = self._make_wavefront_task( - request, phase, token_budget, depth=depth, start=start - ) - if scheduled_request is None: - continue - frontier = scheduled_request.token_start + scheduled_request.token_count - if not self._ensure_execution_capacity(request, frontier, selected=selected): - blocked.add(request_id) - continue - tasks.append((owner_index, scheduled_request)) - selected.add(request_id) - added.add(request_id) - token_budget -= scheduled_request.token_count - - tasks.sort(key=lambda item: (item[0], item[1].token_start)) - scheduled_requests.extend( - replace(scheduled_request, phase=phase) for _, scheduled_request in tasks - ) + if start >= end: + continue + scheduled_request = self._make_wavefront_task( + request, + phase, + token_budget, + depth=(None if phase == RequestPhase.PREFILLING_PRELUDE else depth), + start=start, + ) + if scheduled_request is None or not accept_task(request, scheduled_request): + continue + if phase == RequestPhase.PREFILLING_PRELUDE: + prelude_cursors[request_id] += scheduled_request.token_count + tasks.append((owner_index, scheduled_request)) + tasks.sort(key=lambda item: (item[0], item[1].token_start)) + scheduled_requests.extend( + replace(scheduled_request, phase=phase) for _, scheduled_request in tasks + ) for request in candidates: # Growing another candidate may have suspended this request and put @@ -1158,6 +1150,61 @@ def _wavefront_prefill_groups(self, batch: SchedulerOutput): groups.setdefault((request.request_id, phase), []).append(item) return groups + def _apply_prefill_phase_progress( + self, + scheduled_request: ScheduledRequest, + phase: RequestPhase, + wavefront_groups: dict[tuple[str, RequestPhase], list[ScheduledRequest]], + update: SchedulerUpdate, + *, + loop_count: int, + request_events: dict[str, object] | None = None, + fallback_event: object | None = None, + ) -> bool: + """Apply prefill progress shared by submission and output updates. + + The caller filters stale or finished requests before calling this + helper. Return ``True`` when prefill owns the row, including duplicate + rows consumed by wavefront grouping; return ``False`` for phase-specific + logic. + """ + request = scheduled_request.request + if self.wavefront_chunked_prefill and phase in ( + RequestPhase.PREFILLING_PRELUDE, + RequestPhase.PREFILLING_RECURRENT, + ): + grouped = wavefront_groups.pop((request.request_id, phase), None) + if grouped is not None: + event = ( + request_events.get(request.request_id) + if request_events is not None + else fallback_event + ) + representative, progress = self._apply_wavefront_prefill_batch_progress( + phase, grouped, loop_count=loop_count, event=event + ) + update.prefill.append((representative, phase, progress)) + return True + if phase in ( + RequestPhase.PREFILL, + RequestPhase.PREFILLING_PRELUDE, + RequestPhase.PREFILLING_RECURRENT, + ): + event = ( + request_events.get(request.request_id) + if request_events is not None + else fallback_event + ) + progress = self._apply_prefill_progress( + scheduled_request, + phase, + loop_count=loop_count, + event=event, + ) + update.prefill.append((scheduled_request, phase, progress)) + return True + return False + def _should_exit(self, request: Request) -> bool: """Evaluate the policy after recording the completed loop's actual hazard.""" params = request.sampling_params @@ -1214,6 +1261,7 @@ def update_from_output( result = ModelRunnerOutput() update = SchedulerUpdate() wavefront_groups = self._wavefront_prefill_groups(batch) + loop_count = self._batch_loop_count(batch) for row, scheduled_request in enumerate(batch.requests): request = scheduled_request.request # Ignore results for a request that is no longer registered: it may @@ -1225,40 +1273,17 @@ def update_from_output( ): continue phase = batch.phase_for(scheduled_request) - if self.wavefront_chunked_prefill and phase in ( - RequestPhase.PREFILLING_PRELUDE, - RequestPhase.PREFILLING_RECURRENT, + if self._apply_prefill_phase_progress( + scheduled_request, + phase, + wavefront_groups, + update, + loop_count=loop_count, + request_events=request_events, + fallback_event=result.completion, ): - grouped = wavefront_groups.pop((request.request_id, phase), None) - if grouped is not None: - event = ( - request_events.get(request.request_id) - if request_events is not None - else result.completion - ) - representative, progress = self._apply_wavefront_prefill_batch_progress( - phase, grouped, loop_count=self._batch_loop_count(batch), event=event - ) - update.prefill.append((representative, phase, progress)) continue - if phase in ( - RequestPhase.PREFILL, - RequestPhase.PREFILLING_PRELUDE, - RequestPhase.PREFILLING_RECURRENT, - ): - event = ( - request_events.get(request.request_id) - if request_events is not None - else result.completion - ) - progress = self._apply_prefill_progress( - scheduled_request, - phase, - loop_count=self._batch_loop_count(batch), - event=event, - ) - update.prefill.append((scheduled_request, phase, progress)) - elif phase == RequestPhase.PRELUDE: + if phase == RequestPhase.PRELUDE: request.loops_done = 0 request.remaining_probability = 1.0 request.pending_exit_depth = None @@ -1356,6 +1381,7 @@ def update_after_submission( """Advance logical submissions without requiring CPU token/score delivery.""" update = SchedulerUpdate() wavefront_groups = self._wavefront_prefill_groups(batch) + loop_count = self._batch_loop_count(batch) for row, scheduled_request in enumerate(batch.requests): request = scheduled_request.request if ( @@ -1364,40 +1390,17 @@ def update_after_submission( ): continue phase = batch.phase_for(scheduled_request) - if self.wavefront_chunked_prefill and phase in ( - RequestPhase.PREFILLING_PRELUDE, - RequestPhase.PREFILLING_RECURRENT, + if self._apply_prefill_phase_progress( + scheduled_request, + phase, + wavefront_groups, + update, + loop_count=loop_count, + request_events=request_events, + fallback_event=completion, ): - grouped = wavefront_groups.pop((request.request_id, phase), None) - if grouped is not None: - event = ( - request_events.get(request.request_id) - if request_events is not None - else completion - ) - representative, progress = self._apply_wavefront_prefill_batch_progress( - phase, grouped, loop_count=self._batch_loop_count(batch), event=event - ) - update.prefill.append((representative, phase, progress)) continue - if phase in ( - RequestPhase.PREFILL, - RequestPhase.PREFILLING_PRELUDE, - RequestPhase.PREFILLING_RECURRENT, - ): - event = ( - request_events.get(request.request_id) - if request_events is not None - else completion - ) - progress = self._apply_prefill_progress( - scheduled_request, - phase, - loop_count=self._batch_loop_count(batch), - event=event, - ) - update.prefill.append((scheduled_request, phase, progress)) - elif phase == RequestPhase.PRELUDE: + if phase == RequestPhase.PRELUDE: request.loops_done = 0 request.remaining_probability = 1.0 request.pending_exit_depth = None From b4a848a61b1ce67aeab604bde095767fb68974ed Mon Sep 17 00:00:00 2001 From: LyxWxj Date: Thu, 29 Oct 2026 12:31:54 +0800 Subject: [PATCH 2/6] refactor(scheduler): share split prefill frontier progress Signed-off-by: LyxWxj --- tests/test_split_prefill.py | 55 +++++++++++++++++++++++ vllm_rlt/core/scheduler.py | 89 +++++++++++++++++++++++-------------- 2 files changed, 110 insertions(+), 34 deletions(-) diff --git a/tests/test_split_prefill.py b/tests/test_split_prefill.py index 3c36fb5d..1080eaef 100644 --- a/tests/test_split_prefill.py +++ b/tests/test_split_prefill.py @@ -727,6 +727,61 @@ def test_engine_split_prefill_transition_advances_the_same_frontier_contract(): assert request.phase is RequestPhase.CODA +def test_split_prefill_progress_advances_multiple_recurrent_depths_per_submission(): + from unittest.mock import Mock + + from vllm_rlt.core.scheduler import ScheduledRequest, Scheduler + + cache_manager = Mock() + scheduler = Scheduler( + SchedulerConfig(mode="mixed"), + cache_manager, + total_ut_steps=TOTAL_DEPTHS, + ) + request = Request( + "a", + [1, 2, 3, 4], + SamplingParams(), + prefill_depth_ranges=Request.prefill_boundaries(4, TOTAL_DEPTHS), + ) + staged = ScheduledRequest(request, token_start=0, token_count=2, chunk_start=0) + + progress = scheduler._apply_prefill_progress( + staged, + RequestPhase.PREFILLING_PRELUDE, + loop_count=0, + ) + + assert (progress.chunk_complete, progress.prompt_complete) == (False, False) + assert request.prefill_depth_ranges == [[0, 4], [0, 0], [0, 0], [0, 0]] + scheduler.queues[RequestPhase.PREFILLING_RECURRENT].clear() + + progress = scheduler._apply_prefill_progress( + ScheduledRequest(request, token_start=0, token_count=2, depth=0, chunk_start=0), + RequestPhase.PREFILLING_RECURRENT, + loop_count=2, + ) + + assert (progress.chunk_complete, progress.prompt_complete) == (False, False) + assert request.prefill_depth_ranges == [[2, 4], [2, 2], [0, 2], [0, 0]] + cache_manager.publish_prefix.assert_not_called() + scheduler.queues[RequestPhase.PREFILLING_RECURRENT].clear() + event = object() + + progress = scheduler._apply_prefill_progress( + ScheduledRequest(request, token_start=0, token_count=2, depth=2, chunk_start=0), + RequestPhase.PREFILLING_RECURRENT, + loop_count=2, + event=event, + ) + + assert (progress.chunk_complete, progress.prompt_complete) == (True, False) + assert request.prefill_depth_ranges == [[2, 4], [2, 2], [2, 2], [2, 2]] + assert request.prefill_reclaimable_tokens == 2 + assert request.phase is RequestPhase.PREFILLING_PRELUDE + cache_manager.publish_prefix.assert_called_once_with("a", [1, 2, 3, 4], 2, event) + + @pytest.mark.parametrize( ("chunk_size", "token_budget", "prompt"), [(1, 3, [1]), (2, 3, [1, 2, 3, 4, 5]), (3, 4, [5, 4, 3, 2, 1, 6, 7])], diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index 9d5cd804..22d01eb7 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -955,11 +955,25 @@ def _prompt_destination(self, request: Request) -> RequestPhase: exporting = self.exports_request is not None and self.exports_request(request) return RequestPhase.TRANSFERRING if exporting else RequestPhase.CODA - def _apply_split_prefill_progress( + def _advance_prefill_depth_frontier( self, request: Request, - phase: RequestPhase, + *, + depth: int, + token_end: int, + ) -> bool: + """Advance one submitted prompt span through one recurrent depth.""" + request.prefill_depth_ranges[depth][0] = token_end + if depth + 1 < self.total_ut_steps: + request.prefill_depth_ranges[depth + 1][1] = token_end + return False + request.num_prefilled_tokens = token_end + return True + + def _apply_split_prefill_progress( + self, scheduled_request: ScheduledRequest, + phase: RequestPhase, *, loop_count: int | None = None, event=None, @@ -969,6 +983,7 @@ def _apply_split_prefill_progress( The device work may still be pending. The event orders prefix publication and consumers with that work; this transition does not wait for it. """ + request = scheduled_request.request total_depths = self.total_ut_steps start = scheduled_request.token_start end = start + scheduled_request.token_count @@ -990,19 +1005,19 @@ def _apply_split_prefill_progress( boundaries = request.prefill_depth_ranges if depth + loop_count > total_depths: raise ValueError("prefill execution range exceeds model loop depth") + final_depth_reached = False for offset in range(loop_count): current_depth = depth + offset if phase == RequestPhase.PREFILLING_RECURRENT and ( boundaries[current_depth][0] != start or boundaries[current_depth][1] < end ): raise RuntimeError("prefill execution range crosses an incomplete depth frontier") - boundaries[current_depth][0] = end - if current_depth + 1 < total_depths: - boundaries[current_depth + 1][1] = end - else: - request.num_prefilled_tokens = end + final_depth_reached = self._advance_prefill_depth_frontier( + request, + depth=current_depth, + token_end=end, + ) - final_depth_reached = depth + loop_count == total_depths chunk_complete = final_depth_reached and end == request.num_preluded_tokens if not final_depth_reached or not chunk_complete: self.enqueue(request, RequestPhase.PREFILLING_RECURRENT) @@ -1033,22 +1048,24 @@ def _apply_prefill_progress( event=None, ) -> PrefillProgress: """One prompt frontier update for local and remotely consumed outputs.""" - request = scheduled_request.request - if phase == RequestPhase.PREFILL: - request.num_prefilled_tokens += scheduled_request.token_count - final = request.num_prefilled_tokens == len(request.prompt_token_ids) - self.cache_manager.publish_prefix( - request.request_id, request.prompt_token_ids, request.num_prefilled_tokens, event - ) - if final: - request.loops_done = self.total_ut_steps - self.enqueue( - request, self._prompt_destination(request) if final else RequestPhase.PREFILL + if phase != RequestPhase.PREFILL: + return self._apply_split_prefill_progress( + scheduled_request, + phase, + loop_count=loop_count, + event=event, ) - return PrefillProgress(True, final) - return self._apply_split_prefill_progress( - request, phase, scheduled_request, loop_count=loop_count, event=event + + request = scheduled_request.request + request.num_prefilled_tokens += scheduled_request.token_count + final = request.num_prefilled_tokens == len(request.prompt_token_ids) + self.cache_manager.publish_prefix( + request.request_id, request.prompt_token_ids, request.num_prefilled_tokens, event ) + if final: + request.loops_done = self.total_ut_steps + self.enqueue(request, self._prompt_destination(request) if final else RequestPhase.PREFILL) + return PrefillProgress(True, final) def _apply_wavefront_prefill_batch_progress( self, @@ -1073,12 +1090,16 @@ def _apply_wavefront_prefill_batch_progress( transfer_start = request.prefill_reclaimable_tokens if phase == RequestPhase.PREFILLING_PRELUDE: - frontier = request.num_preluded_tokens - for item in scheduled_requests: - if item.token_start != frontier: - raise RuntimeError("wavefront prelude contains a prompt-range gap") - frontier += item.token_count - request.num_preluded_tokens = frontier + first = scheduled_requests[0] + if first.token_start != request.num_preluded_tokens: + raise RuntimeError("wavefront prelude contains a prompt-range gap") + if any( + previous.token_start + previous.token_count != current.token_start + for previous, current in zip(scheduled_requests, scheduled_requests[1:]) + ): + raise RuntimeError("wavefront prelude contains a prompt-range gap") + last = scheduled_requests[-1] + request.num_preluded_tokens = last.token_start + last.token_count progress = PrefillProgress(False, False, transfer_start) representative = scheduled_requests[0] next_phase = ( @@ -1093,15 +1114,15 @@ def _apply_wavefront_prefill_batch_progress( depth = item.depth start = item.token_start end = start + item.token_count - boundaries = request.prefill_depth_ranges - if boundaries[depth][0] != start: + if request.prefill_depth_ranges[depth][0] != start: raise RuntimeError("stale wavefront recurrent depth frontier") - boundaries[depth][0] = end - if depth + 1 < self.total_ut_steps: - boundaries[depth + 1][1] = end + if not self._advance_prefill_depth_frontier( + request, + depth=depth, + token_end=end, + ): continue - request.num_prefilled_tokens = end chunk_end = min( len(request.prompt_token_ids), item.chunk_start + self.config.prefill_chunk_size, From e6cf85e8180978f656371287e050a56a978ad012 Mon Sep 17 00:00:00 2001 From: LyxWxj Date: Fri, 30 Oct 2026 11:27:34 +0800 Subject: [PATCH 3/6] fix: avoid redundant prelude_cursors construction Signed-off-by: LyxWxj --- vllm_rlt/core/scheduler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index 22d01eb7..2afc1b3f 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -900,7 +900,7 @@ def accept_task(request: Request, scheduled_request: ScheduledRequest) -> bool: tasks = [] prelude_cursors = { request.request_id: request.num_preluded_tokens for request in candidates - } + } if phase == RequestPhase.PREFILLING_PRELUDE else None for depth in range(total_depths): for owner_index, request in enumerate(candidates): request_id = request.request_id @@ -909,6 +909,7 @@ def accept_task(request: Request, scheduled_request: ScheduledRequest) -> bool: if request_id in blocked or request.phase != phase: continue if phase == RequestPhase.PREFILLING_PRELUDE: + assert(prelude_cursors) start, end = prelude_cursors[request_id], len(request.prompt_token_ids) else: start, end = request.prefill_depth_ranges[depth] @@ -925,6 +926,7 @@ def accept_task(request: Request, scheduled_request: ScheduledRequest) -> bool: if scheduled_request is None or not accept_task(request, scheduled_request): continue if phase == RequestPhase.PREFILLING_PRELUDE: + assert(prelude_cursors) prelude_cursors[request_id] += scheduled_request.token_count tasks.append((owner_index, scheduled_request)) tasks.sort(key=lambda item: (item[0], item[1].token_start)) From ddc9aa97eb73c7d8a1273a5f9e8faffac39d7aad Mon Sep 17 00:00:00 2001 From: LyxWxj Date: Fri, 30 Oct 2026 11:32:11 +0800 Subject: [PATCH 4/6] fix: avoid redundant prelude_cursors construction Signed-off-by: LyxWxj --- vllm_rlt/core/scheduler.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index 2afc1b3f..1a8d035a 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -898,9 +898,14 @@ def accept_task(request: Request, scheduled_request: ScheduledRequest) -> bool: return True tasks = [] - prelude_cursors = { - request.request_id: request.num_preluded_tokens for request in candidates - } if phase == RequestPhase.PREFILLING_PRELUDE else None + prelude_cursors = ( + { + request.request_id: request.num_preluded_tokens + for request in candidates + } + if phase == RequestPhase.PREFILLING_PRELUDE + else None + ) for depth in range(total_depths): for owner_index, request in enumerate(candidates): request_id = request.request_id From fa900c77ccff6cd5170d479d43e26bd702cfabef Mon Sep 17 00:00:00 2001 From: LyxWxj Date: Fri, 30 Oct 2026 11:37:02 +0800 Subject: [PATCH 5/6] refactor(scheduler): restore wavefront prelude frontier validation Validate contiguous PRELUDE spans with a single accumulated frontier and avoid the unnecessary pairwise list slice. Signed-off-by: LyxWxj --- vllm_rlt/core/scheduler.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index 1a8d035a..fc829da1 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -1097,16 +1097,12 @@ def _apply_wavefront_prefill_batch_progress( transfer_start = request.prefill_reclaimable_tokens if phase == RequestPhase.PREFILLING_PRELUDE: - first = scheduled_requests[0] - if first.token_start != request.num_preluded_tokens: - raise RuntimeError("wavefront prelude contains a prompt-range gap") - if any( - previous.token_start + previous.token_count != current.token_start - for previous, current in zip(scheduled_requests, scheduled_requests[1:]) - ): - raise RuntimeError("wavefront prelude contains a prompt-range gap") - last = scheduled_requests[-1] - request.num_preluded_tokens = last.token_start + last.token_count + frontier = request.num_preluded_tokens + for item in scheduled_requests: + if item.token_start != frontier: + raise RuntimeError("wavefront prelude contains a prompt-range gap") + frontier += item.token_count + request.num_preluded_tokens = frontier progress = PrefillProgress(False, False, transfer_start) representative = scheduled_requests[0] next_phase = ( From e95daf9e90381dd6f0dce81ca898f4330233f9be Mon Sep 17 00:00:00 2001 From: LyxWxj Date: Fri, 30 Oct 2026 11:46:52 +0800 Subject: [PATCH 6/6] style(scheduler): format prelude cursors and clarify assertions Signed-off-by: LyxWxj --- vllm_rlt/core/scheduler.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index fc829da1..12aa040c 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -899,11 +899,8 @@ def accept_task(request: Request, scheduled_request: ScheduledRequest) -> bool: tasks = [] prelude_cursors = ( - { - request.request_id: request.num_preluded_tokens - for request in candidates - } - if phase == RequestPhase.PREFILLING_PRELUDE + {request.request_id: request.num_preluded_tokens for request in candidates} + if phase == RequestPhase.PREFILLING_PRELUDE else None ) for depth in range(total_depths): @@ -914,7 +911,7 @@ def accept_task(request: Request, scheduled_request: ScheduledRequest) -> bool: if request_id in blocked or request.phase != phase: continue if phase == RequestPhase.PREFILLING_PRELUDE: - assert(prelude_cursors) + assert prelude_cursors is not None start, end = prelude_cursors[request_id], len(request.prompt_token_ids) else: start, end = request.prefill_depth_ranges[depth] @@ -931,7 +928,7 @@ def accept_task(request: Request, scheduled_request: ScheduledRequest) -> bool: if scheduled_request is None or not accept_task(request, scheduled_request): continue if phase == RequestPhase.PREFILLING_PRELUDE: - assert(prelude_cursors) + assert prelude_cursors is not None prelude_cursors[request_id] += scheduled_request.token_count tasks.append((owner_index, scheduled_request)) tasks.sort(key=lambda item: (item[0], item[1].token_start))