Repository: huggingface/transformers
Pull request: [CB] Small fixes

This PR adds a logging message when infering the behavior of use async and fixes an error when evicting a graph from the graph buffer.

Changed files:
filename: benchmark_v2/benchmark_scripts/continuous_batching_overall.py
patch: @@ -52,9 +52,9 @@ def run_and_parse_cb_example(args: str) -> dict:
     results.append(run_and_parse_cb_example("--samples 100 --attn flash_attention_2"))
     results.append(run_and_parse_cb_example("--samples 100 --attn sdpa"))
 
-    # Benchmark with high number of samples
-    results.append(run_and_parse_cb_example("--samples 500"))
-    # Benchmark with async API
+    # Benchmark with high number of samples and synchronous batching
+    results.append(run_and_parse_cb_example("--samples 500 --no-use-async"))
+    # Benchmark with high number of samples and asynchronous batching
     results.append(run_and_parse_cb_example("--samples 500 --use-async"))
 
     # Benchmark with prefix sharing and compile (best performance, but not reproducible due to compilation)

filename: examples/pytorch/continuous_batching.py
patch: @@ -90,7 +90,7 @@ def batch_generate(
     displayed_samples: int = 0,  # -1: no display, 0: display stats, >0: display inputs and some outputs
     output_file: str | None = None,
     expected_outputs: list[str] | None = None,
-    use_async: bool | None = None,
+    use_async_batching: bool | None = None,
 ) -> tuple[float, float]:
     # Actual batch generation
     if displayed_samples >= 0:
@@ -99,7 +99,7 @@ def batch_generate(
     batch_outputs = model.generate_batch(
         inputs=simple_batch_inputs,
         generation_config=generation_config,
-        use_async=use_async,
+        use_async_batching=use_async_batching,
     )
     end_time_simple = time.time()
     if displayed_samples >= 0:
@@ -177,7 +177,7 @@ def batch_generate(
     parser.add_argument("--matmul-precision", "-mp", type=str, default="high")  # set to "none" to disable
     parser.add_argument("--cuda-graph", "-cg", help="Use cuda graphs", type=str, default=None)
     parser.add_argument("--compile", action="store_true", help="Compile the model using torch.compile")
-    parser.add_argument("--use-async", action="store_true", help="Use the asynchronous API for CB")
+    parser.add_argument("--use-async", action=argparse.BooleanOptionalAction, help="Use asynchronous batching")
     parser.add_argument("--do-sample", action="store_true", help="Activate sampling")
     parser.add_argument("--num-return-sequences", type=int, default=1, help="Number of return sequences")
 
@@ -334,7 +334,7 @@ def batch_generate(
             generation_cfg,
             tokenizer,
             displayed_samples=-1,
-            use_async=args.use_async,
+            use_async_batching=args.use_async,
         )
 
     if args.profile is not None:
@@ -351,7 +351,7 @@ def batch_generate(
             displayed_samples=args.displayed,
             output_file=args.output_file,
             expected_outputs=expected_outputs,
-            use_async=args.use_async,
+            use_async_batching=args.use_async,
         )
     if args.profile is not None:
         filename = args.profile if args.profile.endswith(".json") else args.profile + ".json"

filename: src/transformers/generation/continuous_batching/continuous_api.py
patch: @@ -180,7 +180,7 @@ def _get_new_requests(self) -> None:
                 break
             except Exception as e:
                 logger.error(f"Error processing new request: {e}", exc_info=True)
-                state: RequestState = locals().get("state")  # type:ignore
+                state: RequestState = locals().get("state")
                 if state is not None:
                     self._handle_request_error(e, state)
 
@@ -464,7 +464,7 @@ def _process_logit(
         logits_2d = logits.view(batch_size * seq_len, vocab_size)
         input_ids_2d = batch_data["input_ids"].view(batch_size * seq_len)
         # Process with 2D tensors#
-        processed_logits_2d = logit_processor(input_ids_2d, logits_2d)  # type: ignore[arg-type]
+        processed_logits_2d = logit_processor(input_ids_2d, logits_2d)
         # Reshape back to 3D
         return processed_logits_2d.view(batch_size, seq_len, vocab_size)
 
@@ -559,6 +559,11 @@ def __init__(
         # transfers, and if CUDA graphs are not turned off, because that would mean we are trying to save memory
         else:
             self.use_async_batching = self.use_cuda_graph and not attn_mask_is_needed(self.model.config)
+            logger.info(
+                f"No behavior specified for use_async_batching, choosing {self.use_async_batching = } because CUDA "
+                "graphs are turned on and no attention mask is needed. If you want to save memory, you can disable "
+                "asynchronous batching but it will degrade performance (by ~25% for some workloads)."
+            )
 
         # Padding interval sizes for Q and KV (0 means use defaults)
         self.q_padding_interval_size = (
@@ -975,7 +980,7 @@ def init_continuous_batching(
 
         # Create and return the manager
         return ContinuousBatchingManager(
-            model=self,  # type: ignore
+            model=self,
             generation_config=gen_config,
             manual_eviction=manual_eviction,
             max_queue_size=max_queue_size,

filename: src/transformers/generation/continuous_batching/utils.py
patch: @@ -18,6 +18,8 @@
 
 from transformers.configuration_utils import PretrainedConfig
 
+from .requests import logger
+
 
 class CudaGraphBuffer:
     """A fixed-size dict for CUDA graphs with LRU eviction when full."""
@@ -36,8 +38,9 @@ def get_graph(self, q_len: int, kv_len: int) -> torch.cuda.CUDAGraph | None:
 
     def set_graph(self, q_len: int, kv_len: int, graph: torch.cuda.CUDAGraph) -> None:
         if len(self._storage) >= self.max_size:
-            _, graph = self._storage.popitem(last=False)
-            graph.reset()
+            evicted_key, evicted_graph = self._storage.popitem(last=False)
+            logger.info(f"Evicting graph for {evicted_key = }")
+            evicted_graph.reset()
         self._storage[(q_len, kv_len)] = graph