[CB] refactor left padding removal (#211)

### [CB] refactor left padding removal

- calls the function `reduce_left_padding` at every step (prefill and
decode)
- removes the dependency on cached requests
- adjusts /tests for CB covering that exact case

---------

Signed-off-by: Yannick Schnider <[email protected]>

Author: yannicks1

tests/e2e/test_spyre_cb.py

@@ -701,7 +701,7 @@ def get_params_test_remove_left_padding():
         {
             # Prefill sequence 2
             "step": 42,
-            "tkv": 103,  # TODO expecting 39 for next implementation
+            "tkv": 39,  # left padding reduction: 103 - 64 (block size)
             "waiting": [],
             "running": ["2", "1"],
             "request_outputs": ["2"]

vllm_spyre/v1/worker/spyre_model_runner.py

@@ -581,7 +581,7 @@ def __init__(
         # TO DO: move to InputBatch
         self.req_ids2blocks: dict[str, deque[int]] = {}
         self.req_ids2left_pads: dict[str, int] = {}
-        self.tkv = 0
+        self.tkv: int = 0
         self.free_blocks = deque([i for i in range(NUM_BLOCKS)])
         self.dummy_req_ids2blocks: list[int] = []
 
@@ -739,9 +739,6 @@ def _prepare_decode(
                                         dtype=torch.bool,
                                         device="cpu")
 
-        if envs_spyre.VLLM_SPYRE_RM_PADDED_BLOCKS:
-            self.reduce_left_padding(cached_requests)
-
         for cached_request in cached_requests:
             # TODO: Will this always just be one token ID if there's no spec
             # or jump decoding?
@@ -818,27 +815,30 @@ def _prepare_decode(
             is_prompt=False,
         )
 
-    def reduce_left_padding(self, requests: list[CachedRequestData]) -> None:
+    def reduce_left_padding(self) -> None:
+
+        if len(self.req_ids2left_pads) == 0:
+            return
 
-        min_left_pad = min(
-            [self.req_ids2left_pads[r.req_id] for r in requests])
+        min_left_pad = min(self.req_ids2left_pads.values())
         n_padded_blocks = min_left_pad // self.BLOCK_SIZE
+        offset = n_padded_blocks * self.BLOCK_SIZE
 
-        if n_padded_blocks > 0:
+        if offset > 0:
             logger.debug("Number of removed blocks due to left padding: %d",
                          n_padded_blocks)
 
-            for req in requests:
-                self.req_ids2left_pads[
-                    req.req_id] -= n_padded_blocks * self.BLOCK_SIZE
+            for req_id in self.req_ids2left_pads:
+                self.req_ids2left_pads[req_id] -= offset
 
                 # free blocks
                 for _ in range(n_padded_blocks):
-                    freed_block_id = self.req_ids2blocks[req.req_id].popleft()
+                    freed_block_id = self.req_ids2blocks[req_id].popleft()
+                    logger.debug("Freeing block with id: %s", freed_block_id)
                     self.free_blocks.append(freed_block_id)
 
         # update tkv
-        self.tkv -= n_padded_blocks * self.BLOCK_SIZE
+        self.tkv -= offset
 
         return
 
@@ -905,6 +905,10 @@ def pad_input_ids(
     def prepare_model_input(
             self, scheduler_output: SchedulerOutput) -> ModelForwardInputs:
 
+        # remove left padding if applicable before next prefil/decode step
+        if envs_spyre.VLLM_SPYRE_RM_PADDED_BLOCKS:
+            self.reduce_left_padding()
+
         # NOTE: We assume that all sequences in the group are all prompts or
         # all decodes.
         # Also assuming that new sequences are prefills