Repository: huggingface/transformers
Pull request: [CB] Add warmup feature

This PR adds a warmup phase before generation starts, turned on by default. 
It allows for better diagnostics and a more representative user experience than without warmup, where the cost of wamup is payed during the first request rather than in a phase before generation starts. 
Warmup is skipped if cuda graphs and compile are both turned off.

Changed files:
filename: benchmark_v2/benchmark_scripts/continuous_batching_overall.py
patch: @@ -1,3 +1,5 @@
+import argparse
+import json
 import re
 import subprocess
 from pathlib import Path
@@ -8,10 +10,11 @@
 SCRIPT_LOCATION = (Path(__file__).parent.parent.parent / "examples/pytorch/continuous_batching.py").as_posix()
 COMMON_ARGS = "--log-level WARNING --seed 0 --force-max-length".split()
 ERROR_OUTPUT = {"time_seconds": "X", "num_tokens": "X", "throughput_tok_per_sec": "ERROR"}
+RESULTS_FILE = Path(__file__).parent.parent / "benchmark_results/cb_overall_main.json"
 
 
 def run_and_parse_cb_example(args: str) -> dict:
-    print(f"Benchmarking with args: {args}")
+    print(f"\nBenchmarking with args: {args}")
     output = subprocess.run(
         ["python", SCRIPT_LOCATION] + args.split() + COMMON_ARGS,
         stdout=subprocess.PIPE,
@@ -32,7 +35,42 @@ def run_and_parse_cb_example(args: str) -> dict:
         return {"args": args, **ERROR_OUTPUT}
 
 
+def build_comparison_table(results: list[dict], main_results: list[dict]) -> list[dict]:
+    """Build a table comparing current results against saved main results."""
+    main_by_args = {r["args"]: r for r in main_results}
+    comparison = [
+        {
+            "args": "Arguments",
+            "main_tok_per_sec": "Main (tok/s)",
+            "current_tok_per_sec": "Current (tok/s)",
+            "diff_percent": "Diff (%)",
+        }
+    ]
+    for result in results:
+        main = main_by_args.get(result["args"])
+        main_tp = main["throughput_tok_per_sec"] if main else None
+        current_tp = result["throughput_tok_per_sec"]
+        if isinstance(main_tp, (int, float)) and isinstance(current_tp, (int, float)):
+            diff = (current_tp - main_tp) / main_tp * 100
+            diff_str = f"{diff:+.1f}%"
+        else:
+            diff_str = "N/A"
+        comparison.append(
+            {
+                "args": result["args"],
+                "main_tok_per_sec": main_tp if main_tp is not None else "N/A",
+                "current_tok_per_sec": current_tp,
+                "diff_percent": diff_str,
+            }
+        )
+    return comparison
+
+
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--main", action="store_true", help="Save results as the main baseline to compare against.")
+    args = parser.parse_args()
+
     results = [
         {
             "args": "Arguments",
@@ -72,3 +110,16 @@ def run_and_parse_cb_example(args: str) -> dict:
     # Print results
     print()
     print(tabulate(results, tablefmt="github"))
+
+    # The header row is results[0], data rows are results[1:]
+    data_results = results[1:]
+
+    if args.main:
+        RESULTS_FILE.write_text(json.dumps(data_results, indent=2))
+        print(f"\nResults saved to {RESULTS_FILE}")
+    else:
+        if RESULTS_FILE.exists():
+            main_results = json.loads(RESULTS_FILE.read_text())
+            comparison = build_comparison_table(data_results, main_results)
+            print()
+            print(tabulate(comparison, tablefmt="github"))

filename: examples/pytorch/continuous_batching.py
patch: @@ -14,10 +14,7 @@
 import argparse
 import contextlib
 import json
-import logging
 import os
-import time
-from copy import deepcopy
 from itertools import cycle
 
 import datasets
@@ -30,6 +27,9 @@
 from transformers.generation.continuous_batching.requests import logger
 
 
+# TODO: this file is now mostly used for benchmarking, remodel it in consequence.
+
+
 def generate_without_cb(
     model_id: str, sliding_window: int, attn_impl: str, batched_inputs: list[int], generation_config: GenerationConfig
 ) -> dict[str, str]:
@@ -94,15 +94,14 @@ def batch_generate(
     expected_outputs: list[str] | None = None,
 ) -> tuple[float, float]:
     # Actual batch generation
-    if displayed_samples >= 0:
-        print("--- Running CB Generation Example ---")
-    start_time_simple = time.time()
     batch_outputs = model.generate_batch(
         inputs=simple_batch_inputs,
         generation_config=generation_config,
         continuous_batching_config=cb_config,
     )
-    end_time_simple = time.time()
+    generation_start = min(out.created_time for out in batch_outputs.values())
+    generation_end = max(out.lifespan[1] for out in batch_outputs.values())
+
     if displayed_samples >= 0:
         print("Done with batch generation.")
 
@@ -139,7 +138,7 @@ def batch_generate(
             print(f"Request {i} matches" if matches else f"Request {i} does NOT match!")
 
     # Compute stats and maybe print them
-    gen_time = end_time_simple - start_time_simple
+    gen_time = generation_end - generation_start
     tok_per_sec = token_count / gen_time
     if displayed_samples >= 0:
         print("-" * 20)
@@ -327,19 +326,6 @@ def batch_generate(
         attn = args.attn.replace("|", "_").replace("/", "_")
         args.output_file = f"runs/cb/{attn}_{args.samples}_{args.cuda_graph}.json"
 
-    # Run warmup batch generation if log level is above DEBUG # TODO: understand why warmup incurs a large overhead during cache creation
-    if logger.level > logging.DEBUG:
-        gen_cfg = deepcopy(generation_cfg)
-        gen_cfg.max_new_tokens = min(gen_cfg.max_new_tokens, args.block_size + 1)
-        batch_generate(
-            model,
-            batched_inputs,
-            gen_cfg,
-            cb_config,
-            tokenizer,
-            displayed_samples=-1,
-        )
-
     if args.profile is not None:
         cm = profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True)
     else:

filename: src/transformers/generation/configuration_utils.py
patch: @@ -1762,7 +1762,8 @@ def resolve_sentinel_values(self) -> None:
     def resolve_compile_configs(
         self, fallback_compile_config: CompileConfig | None, is_flash_attn: bool, decode_fast_path_available: bool
     ) -> None:
-        """Resolve if the compile configs for varlen and decode paths, modifying these attributes in place if needed."""
+        """Resolve if the compile configs for varlen and decode paths, modifying these attributes in place if needed.
+        Default config use full compile over regional compile, because the throughput is significantly higher (~15%)"""
         logger_ = logging.get_logger("ContinuousBatchingLogger")
 
         # For each config, priority is: explicit config, default config, fallback config, None

filename: src/transformers/generation/continuous_batching/continuous_api.py
patch: @@ -21,6 +21,7 @@
 from contextlib import contextmanager, nullcontext
 from math import ceil
 from time import perf_counter
+from typing import Any
 
 import torch
 from torch import nn
@@ -38,7 +39,7 @@
 from .input_outputs import ContinuousBatchingAsyncIOs, ContinuousBatchingIOs
 from .requests import GenerationOutput, RequestState, RequestStatus, logger
 from .scheduler import SCHEDULER_MAPPING, FIFOScheduler, Scheduler
-from .utils import attn_mask_is_needed, pad_to_interval
+from .utils import attn_mask_is_needed, create_warmup_future_states, pad_to_interval
 
 
 """
@@ -510,27 +511,26 @@ def _generation_step(self, model: nn.Module, logit_processor: LogitsProcessorLis
                     graph.replay()
             # Otherwise, the graph does not exist, so we create it
             else:
-                # TODO: remove this once we are sure there are no race conditions
-                # compute_stream.wait_stream(torch.cuda.current_stream())
-                # Warmup
-                with torch.cuda.stream(compute_stream):
-                    forward_fn(model, batch_data, logit_processor, carry_over_ids, prev_output_ids, output_ids)
-                # torch.cuda.current_stream().wait_stream(compute_stream)
-                # Capture
-                graph = torch.cuda.CUDAGraph()
-                # Continuous batching can run multiple manager threads concurrently in one process, but PyTorch's
-                # default capture mode ("global") errors on CUDA actions from other threads. This means capture can be
-                # invalidated even when each manager uses a different device. To avoid this, we use "thread_local" mode.
-                with torch.cuda.graph(
-                    graph, stream=compute_stream, pool=self.graph_pool, capture_error_mode="thread_local"
-                ):
-                    forward_fn(model, batch_data, logit_processor, carry_over_ids, prev_output_ids, output_ids)
-                # Store
-                self.inputs_and_outputs.set_graph(graph)
+                args = (model, batch_data, logit_processor, carry_over_ids, prev_output_ids, output_ids)
+                self.capture_graph(forward_fn, compute_stream, *args)
 
         # In any case, we transfer the outputs to the host
         self.inputs_and_outputs.retrieve_device_outputs()
 
+    def capture_graph(self, forward_fn: Any, compute_stream: torch.cuda.Stream, *args) -> None:
+        # Warmup (ensures the right result is computed before capturing the graph)
+        with torch.cuda.stream(compute_stream):
+            forward_fn(*args)
+        # Capture
+        graph = torch.cuda.CUDAGraph()
+        # Continuous batching can run multiple manager threads concurrently in one process, but PyTorch's
+        # default capture mode ("global") errors on CUDA actions from other threads. This means capture can be
+        # invalidated even when each manager uses a different device. To avoid this, we use "thread_local" mode.
+        with torch.cuda.graph(graph, stream=compute_stream, pool=self.graph_pool, capture_error_mode="thread_local"):
+            forward_fn(*args)
+        # Store
+        self.inputs_and_outputs.set_graph(graph)
+
     @traced
     def _forward_process_and_sample(
         self,
@@ -606,6 +606,105 @@ def _sample(self, probs: torch.Tensor, logits_indices: torch.Tensor, output_ids:
             # underlying data. It's just a trick to use the same storage for both tensors.
             output_ids[1, :tokens].copy_(logprobs.view(dtype=torch.int32))
 
+    @torch.inference_mode()
+    def warmup(
+        self,
+        model: nn.Module,
+        logit_processor: LogitsProcessorList,
+        num_query_tokens: int = 0,
+        num_cache_tokens: int = 0,
+    ) -> None:
+        """Pre-capture CUDA graphs (or trigger compile warmup) for varlen and decode paths. In async mode, both IO
+        pairs are warmed up since each has its own graph buffer and static tensors."""
+
+        if not self._pad_inputs:
+            logger.info("CUDA graphs and compile are disabled, skipping warmup.")
+            return None
+
+        num_query_tokens = num_query_tokens if num_query_tokens > 0 else self.max_batch_tokens
+        num_query_tokens = min(num_query_tokens, self.max_batch_tokens)
+        num_cache_tokens = num_cache_tokens if num_cache_tokens > 0 else self.cache.block_size * num_query_tokens
+        num_cache_tokens = min(num_cache_tokens, self.cache.num_blocks * self.cache.block_size)
+
+        num_pages = self.cache.num_blocks * self.cache.block_size
+        compute_stream = self.inputs_and_outputs.compute_stream
+
+        # In async mode, each IO pair has its own graph buffer and static tensors, so we warm up both
+        num_io_pairs = 2 if self.use_async_batching else 1
+
+        for pair_idx in range(num_io_pairs):
+            if self.use_async_batching:
+                self.inputs_and_outputs.current_pair = pair_idx
+                logger.info(f"Warming up IO pair {pair_idx + 1}/2...")
+
+            # --- Varlen path ---
+            padded_q = pad_to_interval(num_query_tokens, self.q_padding_interval_size, self.max_batch_tokens)
+            padded_kv = pad_to_interval(num_cache_tokens + num_query_tokens, self.kv_padding_interval_size, num_pages)
+            logger.info(f"Warming up varlen path ({padded_q} Q tokens, {padded_kv} KV tokens)...")
+
+            future_states = create_warmup_future_states(
+                1, RequestStatus.PREFILLING, num_query_tokens, num_cache_tokens, self.cache
+            )
+            try:
+                start = perf_counter()
+                self.inputs_and_outputs.prepare_batch_tensors(future_states, False, padded_q, padded_kv - padded_q)
+                batch_data = self.inputs_and_outputs.get_model_kwargs(use_padding=True)
+                carry_over_ids, prev_output_ids, output_ids = self.inputs_and_outputs.get_cb_kwargs()
+                forward_fn = self._compiled_varlen or self._forward_process_and_sample
+                forward_fn_args = (model, batch_data, logit_processor, carry_over_ids, prev_output_ids, output_ids)
+                if self.use_cuda_graph:
+                    self.capture_graph(forward_fn, compute_stream, *forward_fn_args)
+                else:
+                    with torch.cuda.stream(compute_stream):
+                        forward_fn(*forward_fn_args)
+                logger.info(f"Varlen warmup completed in {perf_counter() - start:.2f}s")
+            except Exception as e:
+                logger.warning(f"Failed to warm up varlen path: {e}")
+            finally:
+                for fs in future_states:
+                    self.cache.free_blocks(fs.state.request_id)
+
+            # Exit here if the decode fast path is not available
+            if self.cache.max_blocks_per_request == 0:
+                continue
+
+            # --- Decode fast path ---
+            logger.info("Warming up decode fast path...")
+            q_interval = self.q_padding_interval_size  # shorthand to avoid overly long lines
+            decode_graphs = 0
+            start = perf_counter()
+            # If N requests reach decoding stage, then the number of query tokens is going to start at N and decrease to
+            # 0 as all request finish. So we warmup for all intervals between 0 and N.
+            for num_requests in range(q_interval, num_query_tokens + q_interval, q_interval):
+                future_states = create_warmup_future_states(
+                    num_requests, RequestStatus.DECODING, 1, self.cache.block_size, self.cache
+                )
+                if not future_states:
+                    continue
+                try:
+                    padded_q = pad_to_interval(len(future_states), q_interval, self.max_batch_tokens)
+                    self.inputs_and_outputs.prepare_batch_tensors(future_states, True, padded_q, 0)
+                    batch_data = self.inputs_and_outputs.get_model_kwargs(use_padding=True)
+                    carry_over_ids, prev_output_ids, output_ids = self.inputs_and_outputs.get_cb_kwargs()
+                    forward_fn = self._compiled_decode or self._forward_process_and_sample
+                    forward_fn_args = (model, batch_data, logit_processor, carry_over_ids, prev_output_ids, output_ids)
+                    if self.use_cuda_graph:
+                        self.capture_graph(forward_fn, compute_stream, *forward_fn_args)
+                    else:
+                        with torch.cuda.stream(compute_stream):
+                            forward_fn(*forward_fn_args)
+                    decode_graphs += 1
+                except Exception as e:
+                    logger.warning(f"Failed to warm up decode path for {num_requests} requests: {e}")
+                finally:
+                    for fs in future_states:
+                        self.cache.free_blocks(fs.state.request_id)
+            logger.info(f"Decode warmup completed ({decode_graphs} graphs) in {perf_counter() - start:.2f}s.")
+
+        # If using async batching, reset to pair 0 for the generation loop
+        if self.use_async_batching:
+            self.inputs_and_outputs.current_pair = 0
+
 
 # Manager Class (User Interface)
 @attach_tracer()
@@ -639,6 +738,7 @@ def __init__(
         self.model = model.eval()
         self.generation_config = generation_config
         self.continuous_batching_config = continuous_batching_config
+        self.warmed_up = False  # Set to True after warmup is completed. Usefull for persistent managers.
         # This is an approximation until the cache is created: it will infer the correct value in cache.__init__
         self._use_prefix_sharing = self.continuous_batching_config.allow_block_sharing
 
@@ -686,6 +786,14 @@ def is_running(self) -> bool:
         """Check if the background generation thread is running."""
         return self._generation_thread is not None and self._generation_thread.is_alive()
 
+    def warmup(self, num_query_tokens: int = 0, num_cache_tokens: int = 0) -> None:
+        """Pre-capture CUDA graphs for varlen and decode paths by running dummy batches. Initializes the batch
+        processor if not already done."""
+        if self.batch_processor is None:
+            self.batch_processor = self._create_batch_processor()
+        self.batch_processor.warmup(self.model, self.logit_processor, num_query_tokens, num_cache_tokens)
+        self.warmed_up = True
+
     # NOTE: don't forget to update `continuous_batching_context_manager` when changing this method's definition
     def stop(self, block: bool = True, timeout: float | None = None, keep_for_next_session: bool = False) -> None:
         """Signal the background thread to stop.
@@ -924,6 +1032,7 @@ def _create_batch_processor(self) -> ContinuousBatchProcessor:
         )
         return batch_processor
 
+    @torch.inference_mode()
     def _run_generation_loop(self) -> None:
         """Main processing loop running in the background thread."""
         try:
@@ -1006,6 +1115,7 @@ class ContinuousMixin:
 
     generation_config: GenerationConfig
 
+    @torch.inference_mode()
     def init_continuous_batching(
         self,
         generation_config: GenerationConfig | None = None,
@@ -1066,26 +1176,34 @@ def destroy_cached_continuous_batching_manager(self) -> None:
             delattr(self, "_cached_continuous_batching_manager")
 
     @contextmanager
+    @torch.inference_mode()
     def continuous_batching_context_manager(
         self,
         generation_config: GenerationConfig | None = None,
         block: bool = True,
         timeout: float | None = None,
         continuous_batching_config: ContinuousBatchingConfig | None = None,
         persistent_manager: bool = False,
+        warmup_requests: int | None = 0,
         **deprecated_kwargs,
     ) -> Generator[ContinuousBatchingManager]:
-        """A context manager to safely use the continuous batching manager. Arguments are similars to the ones of
-        `init_continuous_batching`, expect for:
+        """A context manager to safely use the continuous batching manager. Arguments are similar to the ones of
+        `init_continuous_batching`, except for:
             - block: whether to block the thread when stopping the manager. Default is True.
             - timeout: maximum time to wait for the thread to stop. Default is None (no timeout).
-            - persistent_manager: whether to persist the manager after the context manager is exited. Default is False.
+            - warmup_query_tokens: the number of expected requests for which to warmup. 0 is auto, None is no warmup.
         """
         manager = self.init_continuous_batching(
             generation_config=generation_config,
             continuous_batching_config=continuous_batching_config,
             **deprecated_kwargs,
         )
+        if not (warmup_requests is None or manager.warmed_up):
+            # Warmup is long (~30 sec): best to signal the user it's happening than let them think the manager is stuck
+            logger.warning("Warming up for coninuous batching...")
+            start = perf_counter()
+            manager.warmup(num_query_tokens=warmup_requests, num_cache_tokens=0)
+            logger.warning(f"Warming up completed in {perf_counter() - start:.2f}s.")
         manager.start()
         try:
             yield manager
@@ -1105,6 +1223,7 @@ def generate_batch(
         record_timestamps: bool = False,
         progress_bar: bool = True,
         persistent_manager: bool = False,
+        warmup: bool = True,
         **kwargs,
     ) -> dict[str, GenerationOutput]:
         """Generate sequences for a batch of prompts using continuous batching.
@@ -1116,6 +1235,7 @@ def generate_batch(
             record_timestamps: If set to true, the requests will have a timestamp for each token generated
             progress_bar: If set to true, a progress bar will be displayed
             persistent_manager: whether to persist the manager after the generation is finished. Default is False.
+            warmup: whether to pre-capture CUDA graphs before processing requests. Default is True.
             **kwargs: Additional generation parameters. Only max_new_tokens is used, but other deprecated arguments
                 are extracted and passed to the continuous_batching_config object.
         Returns:
@@ -1144,7 +1264,7 @@ def generate_batch(
         for depr_key in deprecated_keys:
             if depr_key in kwargs:
                 deprecated_kwargs[depr_key] = kwargs.pop(depr_key)
-        # Extract max_new_tokens from kwargs, as it's the only expected kwarg
+        # Extract max_new_tokens from kwargs because it's the only expected kwarg
         max_new_tokens = kwargs.pop("max_new_tokens", None)
 
         # Compute the total number of requests
@@ -1159,6 +1279,7 @@ def generate_batch(
             block=True,
             timeout=5,
             persistent_manager=persistent_manager,
+            warmup_requests=len(inputs) if warmup else None,
             **deprecated_kwargs,
         )
         logging_cm = logging_redirect_tqdm([logger])

filename: src/transformers/generation/continuous_batching/utils.py
patch: @@ -13,12 +13,13 @@
 # limitations under the License.
 from collections import OrderedDict
 from math import ceil
+from typing import Any
 
 import torch
 
 from transformers.configuration_utils import PretrainedConfig
 
-from .requests import logger
+from .requests import FutureRequestState, RequestState, RequestStatus, logger
 
 
 class CudaGraphBuffer:
@@ -158,3 +159,30 @@ def build_attention_mask(
             masked += torch.tril(minus_inf, diagonal=sliding_diagonal)
         # Replace in attention mask
         attention_mask[..., query_range, key_range] = masked
+
+
+def create_warmup_future_states(
+    num: int,
+    status: RequestStatus,
+    num_query_tokens: int,
+    num_cache_tokens: int,
+    cache: Any,  # not annotated to avoid circular import
+) -> list[FutureRequestState]:
+    """An utility function to create a list of FutureRequestStates for the warmup of CB."""
+    # Setup
+    request_ids = [f"__warmup_{status.name}_{i}__" for i in range(num)]
+    total_tokens = num_query_tokens + num_cache_tokens
+    blocks_needed = ceil(total_tokens / cache.block_size)
+    # Main loop
+    future_states = []
+    for req_id in request_ids:
+        state = RequestState(request_id=req_id, initial_tokens=[0] * total_tokens, max_new_tokens=1)
+        state._status = status  # bypass the property setter to avoid the lifecycle side effects
+        state.tokens_to_process = [0] * num_query_tokens
+        state.position_offset = num_cache_tokens
+        # Stop if allocation fails for any request
+        allocated = cache.allocate_blocks(blocks_needed, state.request_id, 0)
+        if allocated is None:
+            return future_states
+        future_states.append(FutureRequestState(state, has_new_token=True, complete_blocks=0))
+    return future_states