Repository: vllm-project/vllm
Pull request: Fix auto prefix bug
Resolves #3193

Fixes a bug that occurs when the entire prompt has already been computed. If the entire prompt is marked as computed, the model runner will attempt to create zero-sized tensors for non-computed prompt tokens. This throws an exception in torch.arrange.

This fix is somewhat of a band-aid because it just makes sure that the last block is never considered computed. This will result in a small amount of unnecessary computation when running the same prompt multiple times.

A more comprehensive fix would update the model runner to behave correctly when it detects that all blocks coming in have already been computed. This can be done in a subsequent change.


Problem statement:
Automatic Prefix Caching Bug
If I enable automatic prefix caching, it occasionally crashes.

```
Future exception was never retrieved
future: <Future finished exception=RuntimeError('step must be nonzero')>
Traceback (most recent call last):
File "/root/vllm/vllm/engine/async_llm_engine.py", line 29, in _raise_exception_on_finish
task.result()
  File "/root/vllm/vllm/engine/async_llm_engine.py", line 412, in run_engine_loop
    has_requests_in_progress = await self.engine_step()
  File "/root/vllm/vllm/engine/async_llm_engine.py", line 391, in engine_step
    request_outputs = await self.engine.step_async()
  File "/root/vllm/vllm/engine/async_llm_engine.py", line 189, in step_async
    all_outputs = await self._run_workers_async(
  File "/root/vllm/vllm/engine/async_llm_engine.py", line 274, in _run_workers_async
    all_outputs = await asyncio.gather(*coros)
  File "/root/miniconda3/envs/vllm/lib/python3.10/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/vllm/vllm/worker/worker.py", line 223, in execute_model
    output = self.model_runner.execute_model(seq_group_metadata_list,
  File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/vllm/vllm/worker/model_runner.py", line 575, in execute_model
    lora_mapping) = self.prepare_input_tensors(seq_group_metadata_list)
  File "/root/vllm/vllm/worker/model_runner.py", line 494, in prepare_input_tensors
    lora_requests) = self._prepare_prompt(seq_group_metadata_list)
File "/root/vllm/vllm/worker/model_runner.py", line 243, in _prepare_prompt
start_loc_tensor = torch.arange(0,
RuntimeError: step must be nonzero

Exception in callback functools.partial(<function _raise_exception_on_finish at 0x7f87f65c35b0>, request_tracker=<vllm.engine.async_llm_engine.RequestTracker object at 0x7f87ec4e3fd0>)
handle: <Handle functools.partial(<function _raise_exception_on_finish at 0x7f87f65c35b0>, request_tracker=<vllm.engine.async_llm_engine.RequestTracker object at 0x7f87ec4e3fd0>)>
Traceback (most recent call last):
  File "/root/vllm/vllm/engine/async_llm_engine.py", line 29, in _raise_exception_on_finish
    task.result()
  File "/root/vllm/vllm/engine/async_llm_engine.py", line 412, in run_engine_loop
    has_requests_in_progress = await self.engine_step()
  File "/root/vllm/vllm/engine/async_llm_engine.py", line 391, in engine_step
    request_outputs = await self.engine.step_async()
  File "/root/vllm/vllm/engine/async_llm_engine.py", line 189, in step_async
    all_outputs = await self._run_workers_async(
  File "/root/vllm/vllm/engine/async_llm_engine.py", line 274, in _run_workers_async
    all_outputs = await asyncio.gather(*coros)
  File "/root/miniconda3/envs/vllm/lib/python3.10/asyncio/tasks.py", line 650, in _wrap_awaitable
    return (yield from awaitable.__await__())
ray.exceptions.RayTaskError(KeyError): ray::RayWorkerVllm.execute_method() (pid=1030270, ip=0.0.0.0, actor_id=be1ed7b0fca5fd6227e71c0101000000, repr=<vllm.engine.ray_utils.RayWorkerVllm object at 0x7f5f2f9ad630>)
  File "/root/vllm/vllm/engine/ray_utils.py", line 37, in execute_method
    return executor(*args, **kwargs)
  File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/vllm/vllm/worker/worker.py", line 212, in execute_model
    num_seq_groups = data["num_seq_groups"]
KeyError: 'num_seq_groups'

The above exception was the direct cause of the following exception:


Traceback (most recent call last):
  File "uvloop/cbhandles.pyx", line 63, in uvloop.loop.Handle._run
  File "/root/vllm/vllm/engine/async_llm_engine.py", line 38, in _raise_exception_on_finish
    raise exc
  File "/root/vllm/vllm/engine/async_llm_engine.py", line 33, in _raise_exception_on_finish
    raise AsyncEngineDeadError(
vllm.engine.async_llm_engine.AsyncEngineDeadError: Task finished unexpectedly. This should never happen! Please open an issue on Github. See stack trace above for the actual cause.
INFO 03-04 20:37:48 async_llm_engine.py:133] Aborted request cmpl-7edf10b340a74b3e8c7c2e07325ae5c6.
ERROR:    Exception in ASGI application
Traceback (most recent call last):
  File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/responses.py", line 264, in __call__
    await wrap(partial(self.listen_for_disconnect, receive))
  File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/responses.py", line 260, in wrap
    await func()
  File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/responses.py", line 237, in listen_for_disconnect
    message = await receive()
  File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/uvicorn/protocols/http/httptools_impl.py", line 580, in receive
    await self.message_event.wait()
  File "/root/miniconda3/envs/vllm/lib/python3.10/asyncio/locks.py", line 214, in wait
    await fut
asyncio.exceptions.CancelledError: Cancelled by cancel scope 7f87bc0d52d0

During handling of the above exception, another exception occurred:

  + Exception Group Traceback (most recent call last):
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/uvicorn/protocols/http/httptools_impl.py", line 419, in run_asgi
  |     result = await app(  # type: ignore[func-returns-value]
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/uvicorn/middleware/proxy_headers.py", line 84, in __call__
  |     return await self.app(scope, receive, send)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/fastapi/applications.py", line 1054, in __call__
  |     await super().__call__(scope, receive, send)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/applications.py", line 123, in __call__
  |     await self.middleware_stack(scope, receive, send)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/middleware/errors.py", line 186, in __call__
  |     raise exc
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/middleware/errors.py", line 164, in __call__
  |     await self.app(scope, receive, _send)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/middleware/cors.py", line 83, in __call__
  |     await self.app(scope, receive, send)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/middleware/exceptions.py", line 62, in __call__
  |     await wrap_app_handling_exceptions(self.app, conn)(scope, receive, send)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/_exception_handler.py", line 64, in wrapped_app
  |     raise exc
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/_exception_handler.py", line 53, in wrapped_app
  |     await app(scope, receive, sender)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/routing.py", line 758, in __call__
  |     await self.middleware_stack(scope, receive, send)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/routing.py", line 778, in app
  |     await route.handle(scope, receive, send)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/routing.py", line 299, in handle
  |     await self.app(scope, receive, send)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/routing.py", line 79, in app
  |     await wrap_app_handling_exceptions(app, request)(scope, receive, send)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/_exception_handler.py", line 64, in wrapped_app
  |     raise exc
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/_exception_handler.py", line 53, in wrapped_app
  |     await app(scope, receive, sender)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/routing.py", line 77, in app
  |     await response(scope, receive, send)
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/responses.py", line 257, in __call__
  |     async with anyio.create_task_group() as task_group:
  |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/anyio/_backends/_asyncio.py", line 678, in __aexit__
  |     raise BaseExceptionGroup(
  | exceptiongroup.ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)
  +-+---------------- 1 ----------------
    | Traceback (most recent call last):
    |   File "/root/vllm/vllm/engine/async_llm_engine.py", line 29, in _raise_exception_on_finish
    |     task.result()
    |   File "/root/vllm/vllm/engine/async_llm_engine.py", line 412, in run_engine_loop
    |     has_requests_in_progress = await self.engine_step()
    |   File "/root/vllm/vllm/engine/async_llm_engine.py", line 391, in engine_step
    |     request_outputs = await self.engine.step_async()
    |   File "/root/vllm/vllm/engine/async_llm_engine.py", line 189, in step_async
    |     all_outputs = await self._run_workers_async(
    |   File "/root/vllm/vllm/engine/async_llm_engine.py", line 274, in _run_workers_async
    |     all_outputs = await asyncio.gather(*coros)
    |   File "/root/miniconda3/envs/vllm/lib/python3.10/asyncio/tasks.py", line 650, in _wrap_awaitable
    |     return (yield from awaitable.__await__())
    | ray.exceptions.RayTaskError(KeyError): ray::RayWorkerVllm.execute_method() (pid=1030270, ip=0.0.0.0, actor_id=be1ed7b0fca5fd6227e71c0101000000, repr=<vllm.engine.ray_utils.RayWorkerVllm object at 0x7f5f2f9ad630>)
    |   File "/root/vllm/vllm/engine/ray_utils.py", line 37, in execute_method
    |     return executor(*args, **kwargs)
    |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    |     return func(*args, **kwargs)
    |   File "/root/vllm/vllm/worker/worker.py", line 212, in execute_model
    |     num_seq_groups = data["num_seq_groups"]
    | KeyError: 'num_seq_groups'
    |
    | The above exception was the direct cause of the following exception:
    |
    | Traceback (most recent call last):
    |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/responses.py", line 260, in wrap
    |     await func()
    |   File "/root/miniconda3/envs/vllm/lib/python3.10/site-packages/starlette/responses.py", line 249, in stream_response
    |     async for chunk in self.body_iterator:
    |   File "/root/vllm/vllm/entrypoints/openai/serving_chat.py", line 148, in chat_completion_stream_generator
    |     async for res in result_generator:
    |   File "/root/vllm/vllm/engine/async_llm_engine.py", line 565, in generate
    |     raise e
    |   File "/root/vllm/vllm/engine/async_llm_engine.py", line 559, in generate
    |     async for request_output in stream:
    |   File "/root/vllm/vllm/engine/async_llm_engine.py", line 69, in __anext__
    |     raise result
    |   File "uvloop/cbhandles.pyx", line 63, in uvloop.loop.Handle._run
    |   File "/root/vllm/vllm/engine/async_llm_engine.py", line 38, in _raise_exception_on_finish
    |     raise exc
    |   File "/root/vllm/vllm/engine/async_llm_engine.py", line 33, in _raise_exception_on_finish
    |     raise AsyncEngineDeadError(
    | vllm.engine.async_llm_engine.AsyncEngineDeadError: Task finished unexpectedly. This should never happen! Please open an issue on Github. See stack trace above for the actual cause.
    +------------------------------------

```

vLLM: main branch
Model: openbuddy-deepseek-67b-v18.1-4k-gptq (Marlin Kernel)
GPU: 4 x RTX3090

Hints from issue discussion:
Can confirm similar issues happened to me as well when automatic prefix caching is enabled.
```
Exception in callback functools.partial(<function _raise_exception_on_finish at 0x7f19b986c0d0>, request_tracker=<vllm.engine.async_llm_engine.RequestTracker object at 0x7f19af5db4f0>)
handle: <Handle functools.partial(<function _raise_exception_on_finish at 0x7f19b986c0d0>, request_tracker=<vllm.engine.async_llm_engine.RequestTracker object at 0x7f19af5db4f0>)>
Traceback (most recent call last):
  File "/workspace/vllm/engine/async_llm_engine.py", line 29, in _raise_exception_on_finish
    task.result()
  File "/workspace/vllm/engine/async_llm_engine.py", line 412, in run_engine_loop
    has_requests_in_progress = await self.engine_step()
  File "/workspace/vllm/engine/async_llm_engine.py", line 391, in engine_step
    request_outputs = await self.engine.step_async()
  File "/workspace/vllm/engine/async_llm_engine.py", line 189, in step_async
    all_outputs = await self._run_workers_async(
  File "/workspace/vllm/engine/async_llm_engine.py", line 274, in _run_workers_async
    all_outputs = await asyncio.gather(*coros)
  File "/usr/lib/python3.10/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/workspace/vllm/worker/worker.py", line 223, in execute_model
    output = self.model_runner.execute_model(seq_group_metadata_list,
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/workspace/vllm/worker/model_runner.py", line 575, in execute_model
    lora_mapping) = self.prepare_input_tensors(seq_group_metadata_list)
File "/workspace/vllm/worker/model_runner.py", line 494, in prepare_input_tensors
INFO 03-05 09:16:34 async_llm_engine.py:133] Aborted request cmpl-49aa25f0dba24ec7b00d8ae6a0a102ad.
    lora_requests) = self._prepare_prompt(seq_group_metadata_list)
  File "/workspace/vllm/worker/model_runner.py", line 243, in _prepare_prompt
    start_loc_tensor = torch.arange(0,
RuntimeError: step must be nonzero

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "uvloop/cbhandles.pyx", line 63, in uvloop.loop.Handle._run
  File "/workspace/vllm/engine/async_llm_engine.py", line 38, in _raise_exception_on_finish
    raise exc
  File "/workspace/vllm/engine/async_llm_engine.py", line 33, in _raise_exception_on_finish
    raise AsyncEngineDeadError(
vllm.engine.async_llm_engine.AsyncEngineDeadError: Task finished unexpectedly. This should never happen! Please open an issue on Github. See stack trace above for the actual cause.
```
Model: `mistralai/Mixtral-8x7B-Instruct-v0.1`, 2xA100-80G, cuda graph enabled.
Note: @78 

> Model: openbuddy-deepseek-67b-v18.1-4k-gptq (Marlin Kernel)

This model is not using the marlin kernel

@SageMoore going to take a look

Commit message: format

Patch to review:
diff --git a/tests/engine/test_computed_prefix_blocks.py b/tests/engine/test_computed_prefix_blocks.py
new file mode 100644
index 000000000000..ed35212cc3f1
--- /dev/null
+++ b/tests/engine/test_computed_prefix_blocks.py
@@ -0,0 +1,34 @@
+import pytest
+
+from vllm.engine.arg_utils import EngineArgs
+from vllm.engine.llm_engine import LLMEngine
+from vllm.sampling_params import SamplingParams
+
+
+@pytest.mark.parametrize("model", ["facebook/opt-125m"])
+@pytest.mark.parametrize("block_size", [16])
+def test_computed_prefix_blocks(model: str, block_size: int):
+    # This test checks if we are able to run the engine to completion
+    # without triggering asserts.
+    # We are in a scenario where all blocks from the second request's prompt
+    # are full and already computed when the second request arrives.
+    prompt = (
+        "You are a helpful assistant. How do I build a car from cardboard and "
+        "paper clips? Is there an easy to follow video tutorial available "
+        "online for free?")
+    prompt2 = (
+        " Please recommend to me some resources where I can learn not only to "
+        "handle technical difficulties of building a car, but also "
+        "decoration.")
+
+    engine_args = EngineArgs(model=model,
+                             block_size=block_size,
+                             enable_prefix_caching=True)
+
+    engine = LLMEngine.from_engine_args(engine_args)
+    sampling_params = SamplingParams()
+
+    engine.add_request("0", prompt + prompt2, sampling_params)
+    engine.step()
+    engine.add_request("1", prompt, sampling_params)
+    engine.step()
diff --git a/vllm/core/block_manager.py b/vllm/core/block_manager.py
index daf83827a7e5..b951bb76de7b 100644
--- a/vllm/core/block_manager.py
+++ b/vllm/core/block_manager.py
@@ -1,6 +1,6 @@
 """A block manager that manages token blocks."""
 import enum
-from itertools import count
+from itertools import count, takewhile
 from os.path import commonprefix
 from typing import Dict, List, Optional, Set, Tuple
 
@@ -426,23 +426,29 @@ def access_all_blocks_in_seq(
         for block in block_table:
             block.last_accessed = access_time
 
-    def compute_last_full_block_in_seq(self, seq: Sequence):
+    def compute_full_blocks_in_seq(self, seq: Sequence):
         if seq.seq_id not in self.block_tables:
             return
         max_full_block = seq.get_len() // self.block_size - 1
         block_table = self.block_tables[seq.seq_id]
         if max_full_block == -1:
             return
-        block_table[max_full_block].computed = True
+        for i in reversed(range(max_full_block)):
+            if block_table[i].computed:
+                break
+            block_table[i].computed = True
 
-    def get_all_block_ids_till_computed(self, seq: Sequence) -> List[int]:
+    def get_all_computed_blocks(self, seq: Sequence) -> List[int]:
         if seq.seq_id not in self.block_tables:
             return []
         block_table = self.block_tables[seq.seq_id]
-        for block_idx in reversed(range(len(block_table))):
-            if block_table[block_idx].computed:
-                return [b.block_number for b in block_table[:block_idx + 1]]
-        return []
+        # TODO We exclude the last block to avoid the case where the entire
+        # prompt is cached. This would currently cause erroneous behavior in
+        # worker.
+        return [
+            b.block_number
+            for b in takewhile(lambda b: b.computed, block_table[:-1])
+        ]
 
     def get_common_computed_block_ids(self,
                                       seq_group: SequenceGroup) -> List[int]:
@@ -451,14 +457,12 @@ def get_common_computed_block_ids(self,
             return []
 
         ids_list = [
-            self.get_all_block_ids_till_computed(seq)
+            self.get_all_computed_blocks(seq)
             for seq in iter(seq_group.seqs_dict.values())
         ]
         return commonprefix([ids for ids in ids_list if ids != []])
 
     def mark_blocks_as_computed(self, seq_group: SequenceGroup):
-        # NOTE: We only mark the last full block because with prefix caching,
-        # all blocks until the marked one are guaranteed to be computed.
         if self.enable_caching:
             for seq in seq_group.seqs_dict.values():
-                self.compute_last_full_block_in_seq(seq)
+                self.compute_full_blocks_in_seq(seq)
diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py
index aff8ebc90362..9516d4bdea94 100644
--- a/vllm/worker/model_runner.py
+++ b/vllm/worker/model_runner.py
@@ -209,6 +209,7 @@ def _prepare_prompt(
                 slot_mapping[-1].append(slot)
 
         max_prompt_len = max(subquery_lens)
+        assert max_prompt_len > 0
         input_tokens = _make_tensor_with_pad(input_tokens,
                                              max_prompt_len,
                                              pad=0,


Code review comments:
diff_hunk: @@ -426,23 +426,29 @@ def access_all_blocks_in_seq(
         for block in block_table:
             block.last_accessed = access_time
 
-    def compute_last_full_block_in_seq(self, seq: Sequence):
+    def compute_full_blocks_in_seq(self, seq: Sequence):
         if seq.seq_id not in self.block_tables:
             return
         max_full_block = seq.get_len() // self.block_size - 1
         block_table = self.block_tables[seq.seq_id]
         if max_full_block == -1:
             return
-        block_table[max_full_block].computed = True
+        for i in reversed(range(max_full_block)):
+            if block_table[i].computed:
+                break
+            block_table[i].computed = True
 
-    def get_all_block_ids_till_computed(self, seq: Sequence) -> List[int]:
+    def get_all_computed_blocks(self, seq: Sequence) -> List[int]:
         if seq.seq_id not in self.block_tables:
             return []
         block_table = self.block_tables[seq.seq_id]
-        for block_idx in reversed(range(len(block_table))):
-            if block_table[block_idx].computed:
-                return [b.block_number for b in block_table[:block_idx + 1]]
-        return []
+        # TODO We exclude the last block to avoid the case where the entire
+        # prompt is cached. This would currently cause erroneous behavior in
+        # worker.
path: vllm/core/block_manager.py
text: @user1:
this seems more like a NOTE rather than a TODO, is there future work to consider here such as allowing workers to have zero work?

@author:
Yes, ideally, the workers should be rewritten to handle zero workloads in the future. The current implementation has a small drawback of not being able to read the last block in sequence from cache even if it's there.

@user2:
Let's change this to a NOTE. Even if the whole prefix has been computed before, we still need to recompute on the last token to sample for the first output token. This is unavoidable.