clear prefill metadata

vllm-project · Sep 17, 2024 · 0ffa5fd · 0ffa5fd
1 parent 427ef77
commit 0ffa5fd
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/vllm/worker/multi_step_worker.py b/vllm/worker/multi_step_worker.py
@@ -76,8 +76,9 @@ def _get_driver_input_and_broadcast(
             frozen_model_input = model_input.frozen_model_input
             assert frozen_model_input is not None
             assert frozen_model_input.attn_metadata is not None
-            # clear the cached decode metadata so that it can be recomputed on
-            # the workers
+            # clear the cached metadata so that it can be recomputed on
+            # the workers.
+            frozen_model_input.attn_metadata._cached_prefill_metadata = None
             frozen_model_input.attn_metadata._cached_decode_metadata = None
 
         model_input.is_first_multi_step = is_first_multi_step