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 5, 19:45 UTC
[MRv2]fix: model accuracy regression caused by reusing the stale last_sampled_tokens and draft_tokens (#39833)
Signed-off-by: Zijing Liu <liuzijing2014@gmail.com>
Zijing Liu committed
Apr 21, 2026 at 09:30 UTC
9a6a66f3b837bd3565471dc09ce3e23831e0e3f7
1 file changed
+12
vllm/v1/worker/gpu/states.py
+12
@@ -102,6 +102,18 @@ class RequestState:
102
self.num_computed_prefill_tokens[req_idx] = num_computed_tokens
103
self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens)
104
105
+ if num_computed_tokens > 0 and num_computed_tokens <= prefill_len:
106
+ # For PD disagg or resumed requests: set last_sampled to the last
107
+ # computed token so the first decode step gets the right input_id.
108
+ # For fresh prefill requests (num_computed_tokens == 0) the tensor
109
+ # is not read by combine_sampled_and_draft_tokens so we skip the
110
+ # write. Use a slice assignment rather than scalar indexing so the
111
+ # write is dispatched through fill_ without a host/device sync.
112
+ self.last_sampled_tokens[req_idx : req_idx + 1] = all_token_ids[
113
+ num_computed_tokens - 1
114
+ ]
115
+ self.draft_tokens[req_idx].zero_()
116
+
117
def apply_staged_writes(self) -> None:
118
self.prompt_len.copy_to_uva()
119
self.prefill_len.copy_to_uva()