A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 08:00 UTC
[bug] AsyncScheduler drops first post-resume token after pause_generation + clear_cache (#42117)
Signed-off-by: hao-aaron <ahao@anyscale.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Aaron Hao committed
May 19, 2026 at 01:06 UTC
f34623bf3cac5b33451a761e802c9531e83d1c68
4 files changed
+13
-15
examples/rl/rlhf_async_new_apis.py
-6
@@ -227,9 +227,6 @@ llm_kwargs = dict(
227
attention_backend=ATTN_BACKEND,
228
gpu_memory_utilization=0.75,
229
weight_transfer_config=WeightTransferConfig(backend="nccl"),
230
- # TODO(haosdent): re-enable once #42043 is fixed. Both LLM
231
- # instances must match.
232
- async_scheduling=False,
230
)
231
llm_kwargs.update(rocm_determinism_kwargs)
232
@@ -371,9 +368,6 @@ llm_v2_kwargs = dict(
368
gpu_memory_utilization=0.75,
369
distributed_executor_backend="ray",
370
attention_backend=ATTN_BACKEND,
374
- # TODO(haosdent): re-enable once #42043 is fixed. Both LLM
375
- # instances must match.
376
- async_scheduling=False,
371
)
372
llm_v2_kwargs.update(rocm_determinism_kwargs)
373
vllm/v1/core/sched/async_scheduler.py
+5
-4
@@ -37,10 +37,11 @@ class AsyncScheduler(Scheduler):
37
def _update_request_with_output(
38
self, request: Request, new_token_ids: list[int]
39
) -> tuple[list[int], bool]:
40
- if request.discard_latest_async_tokens:
41
- # If the request is force preempted in reset_prefix_cache, we
42
- # should discard the latest async token.
43
- request.discard_latest_async_tokens = False
40
+ if request.async_tokens_to_discard > 0:
41
+ # The request was force-preempted in reset_prefix_cache; drop one
42
+ # stale in-flight async output frame per call until the counter
43
+ # is drained.
44
+ request.async_tokens_to_discard -= 1
45
return [], False
46
47
status_before_update = request.status
vllm/v1/core/sched/scheduler.py
+7
-3
@@ -1905,10 +1905,14 @@ class Scheduler(SchedulerInterface):
1905
while self.running:
1906
request = self.running.pop()
1907
self._preempt_request(request, timestamp)
1908
- # NOTE(zhuohan): For async scheduling, we need to discard the latest
1909
- # output token on the fly to avoid a redundant repetitive output token.
1908
+ # For async scheduling, any output frames already in flight at
1909
+ # preemption time are now stale and must be discarded when they
1910
+ # return. num_output_placeholders is exactly that count: 0 if
1911
+ # the engine has drained (e.g. pause_generation(keep) waited
1912
+ # for idle), 1 for vanilla async mid-step, or 1 + spec/PP frames
1913
+ # otherwise.
1914
+ request.async_tokens_to_discard = request.num_output_placeholders
1915
request.num_output_placeholders = 0
1911
- request.discard_latest_async_tokens = True
1916
1917
# Clear scheduled request ids cache. Since we are forcing preemption
1918
# + resumption in the same step, we must act as if these requests were
vllm/v1/request.py
+1
-2
@@ -139,8 +139,7 @@ class Request:
139
140
# Used in async scheduling.
141
self.num_output_placeholders = 0
142
- # Used in forced preemption (reset_prefix_cache) with async scheduling.
143
- self.discard_latest_async_tokens = False
142
+ self.async_tokens_to_discard = 0
143
144
self.spec_token_ids: list[int] = []
145
self.num_computed_tokens = 0