Merge pull request vllm-project#8 from robertgshaw2-redhat/avoid-unnatural-free

updated

Author: robertgshaw2-redhat

vllm/v1/core/kv_cache_manager.py

@@ -107,6 +107,7 @@ def get_computed_blocks(
                 - A list of blocks that are computed for the request.
                 - The number of computed tokens.
         """
+
         if not self.enable_caching:
             # Prefix caching is disabled.
             return [], 0
@@ -274,8 +275,6 @@ def allocate_slots(
         # For disaggregated, avoid caching until KVs are recved.
         if skip_cache_blocks:
             assert request.request_id not in self.num_cached_block
-            self.num_cached_block[request.request_id] = len(
-                new_computed_blocks)
             return new_blocks
 
         self.cache_blocks(

vllm/v1/core/sched/scheduler_disagg.py

@@ -196,18 +196,22 @@ def schedule(self) -> SchedulerOutput:
                 # for the requests to arrive.
                 if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
                     if request.request_id in self.finished_recving_kv_req_ids:
-                        assert self.kv_cache_manager.enable_caching
                         # Now that the KVs have been recved, we can cache
                         # them and set num_computed_tokens.
+                        blocks = self.kv_cache_manager.req_to_blocks[
+                            request.request_id]
+                        num_computed_tokens = len(blocks) * self.block_size
                         self.kv_cache_manager.cache_blocks(
                             request,
                             num_tokens=0,
-                            num_computed_tokens=(len(request.all_token_ids) -
-                                                 1))
+                            num_computed_tokens=num_computed_tokens)
+                        assert blocks[-1].block_hash is not None
+
+                        request.num_computed_tokens = num_computed_tokens
+                        request.status = RequestStatus.WAITING
+
                         self.finished_recving_kv_req_ids.remove(
                             request.request_id)
-                        request.status = RequestStatus.WAITING
-                        self.kv_cache_manager.free(request)
                     else:
                         self.waiting.popleft()
                         skipped_waiting_requests.appendleft(request)
@@ -224,10 +228,16 @@ def schedule(self) -> SchedulerOutput:
                         skipped_waiting_requests.appendleft(request)
                         continue
 
-                # Get already-cached tokens.
-                computed_blocks, num_computed_tokens = \
-                    self.kv_cache_manager.get_computed_blocks(
-                        request)
+                # TODO: comment.
+                request_blocks = self.kv_cache_manager.req_to_blocks.get(
+                    request.request_id, None)
+                if request_blocks:
+                    new_computed_blocks = []
+                    num_computed_tokens = len(request_blocks) * self.block_size
+                else:
+                    # Get already-cached tokens.
+                    new_computed_blocks, num_computed_tokens = (
+                        self.kv_cache_manager.get_computed_blocks(request))
 
                 # Get externally-cached tokens if using a KVConnector.
                 num_external_tokens = (
@@ -244,7 +254,7 @@ def schedule(self) -> SchedulerOutput:
                     new_blocks = self.kv_cache_manager.allocate_slots(
                         request,
                         num_external_tokens,
-                        computed_blocks,
+                        new_computed_blocks,
                         skip_cache_blocks=True)
                     if new_blocks is None:
                         # Requests cannot be scheduled
@@ -262,7 +272,7 @@ def schedule(self) -> SchedulerOutput:
                             request,
                             [
                                 b.block_id for b in itertools.chain(
-                                    computed_blocks, new_blocks)
+                                    new_computed_blocks, new_blocks)
                             ],
                             num_external_tokens,
                         )
@@ -274,6 +284,8 @@ def schedule(self) -> SchedulerOutput:
                 # We use `request.num_tokens` instead of
                 # `request.num_prompt_tokens` to consider the resumed request,
                 # which have output tokens.
+                print(f"{request.num_tokens=}")
+                print(f"{num_computed_tokens=}")
                 num_new_tokens = request.num_tokens - num_computed_tokens
                 if (0 < self.scheduler_config.long_prefill_token_threshold <
                         num_new_tokens):
@@ -298,7 +310,7 @@ def schedule(self) -> SchedulerOutput:
                 new_blocks = self.kv_cache_manager.allocate_slots(
                     request,
                     num_new_tokens + num_external_tokens,
-                    computed_blocks,
+                    new_computed_blocks,
                 )
                 if new_blocks is None:
                     # The request cannot be scheduled.
@@ -312,7 +324,7 @@ def schedule(self) -> SchedulerOutput:
                         request,
                         [
                             b.block_id for b in itertools.chain(
-                                computed_blocks, new_blocks)
+                                new_computed_blocks, new_blocks)
                         ],
                         num_external_tokens,
                     )
@@ -335,7 +347,7 @@ def schedule(self) -> SchedulerOutput:
                         f"Invalid request status: {request.status}")
 
                 req_to_new_block_ids[request.request_id] = [
-                    b.block_id for b in computed_blocks + new_blocks
+                    b.block_id for b in new_computed_blocks + new_blocks
                 ]
                 num_scheduled_tokens[request.request_id] = num_new_tokens
                 token_budget -= num_new_tokens