Style

Author: remi-or

src/transformers/generation/continuous_batching/cache_manager.py

@@ -264,7 +264,9 @@ def get_write_indices(self, request_id: str, past_length: int, query_length: int
     def get_seqlens_k(self, request_id: str, past_length: int, query_length: int) -> tuple[str, int]:
         """Returns the attention type of the cache allocator and the key sequence length for the given request_id."""
 
-    def fork_blocks(self, source_request_id: str, dst_request_id: str, block_manager: BlockManager) -> tuple[list[int], list[int]]:
+    def fork_blocks(
+        self, source_request_id: str, dst_request_id: str, block_manager: BlockManager
+    ) -> tuple[list[int], list[int]]:
         """Fork the cache blocks for a given request_id into a new request_id."""
         if source_request_id not in self.block_table:
             raise ValueError(f"No block table found for request {source_request_id}")
@@ -283,6 +285,7 @@ def fork_blocks(self, source_request_id: str, dst_request_id: str, block_manager
         self.block_table[dst_request_id] = forked_blocks
         return source_blocks, forked_blocks
 
+
 class FullAttentionCacheAllocator(CacheAllocator):
     """Cache manager for a group of full attention layers."""
 

src/transformers/generation/continuous_batching/continuous_api.py

@@ -576,7 +576,6 @@ def update_batch(self) -> None:
         for state in self.requests_in_batch:
             # If the request has no remaining prompt ids, it means prefill has already ended or just finished
             if len(state.remaining_prefill_tokens) == 0:
-
                 # If there are no generated tokens yet, it means prefill just ended
                 if state.generated_len() == 0:
                     self.metrics.record_ttft_metric(state.created_time, state.request_id)