[bugfix] Mark draft tokens to rebuilt their embeddings. (#57356)
Signed-off-by: HieDean <799287043@qq.com> Co-authored-by: hiedean <hiedean@tju.edu.cn>
HieDean committed
Sep 17, 2026 at 23:20 UTC
493d4b3c04a439c926df2c4d2fa3cdfabb1a8605
1 file changed
+6
-2
vllm/v1/worker/gpu_model_runner.py
+6
-2
@@ -1858,8 +1858,12 @@ class GPUModelRunner(
1858
total_without_spec = total_num_scheduled_tokens - total_num_spec_tokens
1859
if self.enable_prompt_embeds:
1860
# The multimodal embed path reads is_token_ids.gpu; its .cpu copy is
1861
- # refreshed every step but the async fast paths below only scatter
1862
- # input_ids.gpu, so refresh is_token_ids.gpu here too.
1861
+ # refreshed every step but the async fast path below replaces draft
1862
+ # tokens directly on the GPU. Mark those positions as token IDs before
1863
+ # refreshing the GPU mask so their embeddings are rebuilt. Sampled
1864
+ # tokens are already marked by bookkeeping.
1865
+ if spec_flattened_indices:
1866
+ self.is_token_ids.np[spec_flattened_indices] = True
1867
self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)
1868
if num_common_tokens < total_without_spec:
1869
# If not all requests are decodes from the last iteration,