Repository: huggingface/transformers
Pull request: Make benchmarking lighter: clean-up result files and remove non-needed arguments

This PR simplifies the benchmarking suite by:
- removing the timestamps for each token generation from the result file: we now only keep the time to first token and the (mean) inter-token latency
- removes the SDPA backend argument which is both un-used in practice and introduces a lot of edge cases
- starts the refactoring of `time_generate_batch` and fixes some logic in it (it was not recording the right throughput / results)
- adds a bit more context to loading bars and removes some messages

Changed files:
filename: benchmark_v2/framework/benchmark_config.py
patch: @@ -2,9 +2,10 @@
 import itertools
 import json
 import logging
+from functools import lru_cache
 from typing import Any
 
-from transformers.utils.import_utils import is_flash_attn_2_available
+from transformers.utils.import_utils import is_flash_attn_2_available, is_kernels_available
 
 
 KERNELIZATION_AVAILABLE = False
@@ -18,17 +19,36 @@
 logger = logging.getLogger(__name__)
 
 
+@lru_cache
+def is_fa2_or_kernel_available() -> bool:
+    """Returns True if the flash_attn_2 or a fallback kernel is available"""
+    # Early return if flash_attn_2 is available
+    if is_flash_attn_2_available():
+        return True
+    # Early return if kernels is not available
+    if not is_kernels_available():
+        logger.warning(
+            "flash_attention_2 is not available. kernels is not installed. Benchmarking flash_attention_2 will not "
+            "be possible."
+        )
+        return False
+    # If kernels is available, try to get the flash_attn_2 kernel
+    try:
+        from kernels import get_kernel
+
+        get_kernel("kernels-community/flash-attn")
+    except Exception as _:
+        logger.warning(
+            "flash_attention_2 is not available. kernels is installed, but the flash_attn kernel is not available."
+            "Benchmarking flash_attention_2 will not be possible."
+        )
+        return False
+
+
 class BenchmarkConfig:
     """Configuration for a single benchmark scenario."""
 
-    all_attn_implementations = [
-        ("flash_attention_2", None),
-        ("eager", None),
-        ("sdpa", "math"),
-        ("sdpa", "flash_attention"),
-        ("flex_attention", None),
-    ]
-
+    all_attn_implementations = ["flash_attention_2", "eager", "sdpa", "flex_attention"]
     all_compiled_modes = [None, "default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"]
 
     def __init__(
@@ -41,7 +61,6 @@ def __init__(
         sequence_length: int = 128,
         num_tokens_to_generate: int = 128,
         attn_implementation: str = "eager",
-        sdpa_backend: str | None = None,
         compile_mode: str | None = None,
         compile_options: dict[str, Any] | None = None,
         kernelize: bool = False,
@@ -59,7 +78,6 @@ def __init__(
         self.num_tokens_to_generate = num_tokens_to_generate
         # Generation parameters
         self.attn_implementation = attn_implementation
-        self.sdpa_backend = sdpa_backend
         # Optimization parameters
         self.compile_mode = compile_mode
         self.compile_options = compile_options if compile_options is not None else {}
@@ -75,34 +93,21 @@ def check_validity(self, skip_validity_check: bool = False) -> None:
         if skip_validity_check:
             return
         # Check FA is installed
-        if self.attn_implementation == "flash_attention_2" and not is_flash_attn_2_available():
-            logger.warning(
-                "Flash attention does not support compile mode. Defaulting to SDPA w/ flash attention backend."
-            )
+        is_fa = self.attn_implementation == "flash_attention_2"
+        if is_fa and not is_fa2_or_kernel_available():
+            logger.warning("Flash attention is not available. Defaulting to SDPA.")
             self.attn_implementation = "sdpa"
-            self.sdpa_backend = "flash_attention"
         # Flash attention does not support compile mode, so we turn it off # FIXME: it would be better to support it
-        is_fa = self.attn_implementation == "flash_attention_2"
-        is_fa |= self.attn_implementation == "sdpa" and self.sdpa_backend == "flash_attention"
-        if is_fa:
+        if is_fa and self.compile_mode is not None:
             logger.warning("Flash attention does not support compile mode. Turning off compile mode.")
             self.compile_mode = None
-        # Handle SDPA backend if not determined by the config (needs to be done before skipping duplicates)
-        if self.attn_implementation == "sdpa" and self.sdpa_backend is None:
-            default_backend = "flash_attention"  # FIXME: torch has a _cur_sdpa_kernel_backends but it fails
-            logger.warning(f"No SDPA backend provided, using {default_backend} instead.")
-            self.sdpa_backend = default_backend
+        # Handle continuous batching cases
         if self.continuous_batching:
             if self.attn_implementation == "flex_attention":
                 logger.error(
-                    "disabling continuous batching because of invalid configuration: flex attention is not supported"
+                    "Disabling continuous batching because of invalid configuration: flex attention is not supported."
                 )
                 self.continuous_batching = False
-            elif self.attn_implementation == "sdpa" and self.sdpa_backend is not None:
-                logger.warning(
-                    "when continuous batching is enabled, sdpa_backend must be None because of the attention mask, setting it to None"
-                )
-                self.sdpa_backend = None
 
     @property
     def hash(self) -> str:
@@ -115,7 +120,6 @@ def infer_name(self, compact: bool = True) -> str:
             gpu_monitor_str = "monitored" if self.gpu_monitoring else "unmonitored"
             dimensions_str = f"b{self.batch_size}_s{self.sequence_length}_n{self.num_tokens_to_generate}"
             attn_code = self.attn_implementation
-            attn_code += f"_{self.sdpa_backend}" if self.attn_implementation == "sdpa" else ""
             compile_str = f"compiled_{self.compile_mode}" if self.compile_mode is not None else "uncompiled"
             kernelize_str = "kernelized" if self.kernelize else "unkernelized"
             continuous_batching_str = "cb" if self.continuous_batching else "generate"
@@ -125,7 +129,6 @@ def infer_name(self, compact: bool = True) -> str:
             gpu_monitor_str = ("with" if self.gpu_monitoring else "no") + " GPU monitoring"
             dimensions_str = f"batch size {self.batch_size}, sequence length {self.sequence_length}, {self.num_tokens_to_generate} generated tokens"
             attn_code = f"{self.attn_implementation} attention"
-            attn_code += f" with {self.sdpa_backend} backend" if self.attn_implementation == "sdpa" else ""
             compile_str = "compiled" if self.compile_mode is not None else "not compiled"
             kernelize_str = "kernelized" if self.kernelize else "not kernelized"
             continuous_batching_str = "continuous batching" if self.continuous_batching else "regular generate"
@@ -145,7 +148,6 @@ def to_dict(self) -> dict[str, Any]:
             "sequence_length": self.sequence_length,
             "num_tokens_to_generate": self.num_tokens_to_generate,
             "attn_implementation": self.attn_implementation,
-            "sdpa_backend": self.sdpa_backend,
             "compile_mode": self.compile_mode,
             "compile_options": self.compile_options | {},  # to avoid inplace modification of the original dict
             "kernelize": self.kernelize,
@@ -162,7 +164,6 @@ def from_dict(cls, data: dict[str, Any], skip_validity_check: bool = False) -> "
             sequence_length=data.get("sequence_length", 128),
             num_tokens_to_generate=data.get("num_tokens_to_generate", 128),
             attn_implementation=data.get("attn_implementation", "eager"),
-            sdpa_backend=data.get("sdpa_backend"),
             compile_mode=data.get("compile_mode"),
             compile_options=data.get("compile_options"),
             kernelize=data.get("kernelize", False),
@@ -213,7 +214,7 @@ def get_config_by_level(level: int) -> list[BenchmarkConfig]:
     configs = []
     # Early return if level is greater than 3: we generate all combinations of configs, maybe even w/ all compile modes
     if level >= 3:
-        for attn_implementation, sdpa_backend in BenchmarkConfig.all_attn_implementations:
+        for attn_implementation in BenchmarkConfig.all_attn_implementations:
             # Usually there is not much to gain by compiling with other modes, but we allow it for level 4
             compile_modes = BenchmarkConfig.all_compiled_modes if level >= 4 else [None, "default"]
             for cm in compile_modes:
@@ -222,7 +223,6 @@ def get_config_by_level(level: int) -> list[BenchmarkConfig]:
                         configs.append(
                             BenchmarkConfig(
                                 attn_implementation=attn_implementation,
-                                sdpa_backend=sdpa_backend,
                                 compile_mode=cm,
                                 kernelize=kernelize_on,
                                 continuous_batching=cb_on,

filename: benchmark_v2/framework/benchmark_runner.py
patch: @@ -6,7 +6,6 @@
 import re
 import tempfile
 import time
-from contextlib import nullcontext
 from datetime import datetime
 from queue import Queue
 from typing import Any
@@ -79,24 +78,6 @@ def get_git_revision() -> str:
         return git_hash.readline().strip()
 
 
-def get_sdpa_backend(backend_name: str | None) -> torch.nn.attention.SDPBackend | None:
-    """Get the SDPA backend enum from string name."""
-    if backend_name is None:
-        return None
-
-    try:
-        backend_map = {
-            "math": torch.nn.attention.SDPBackend.MATH,
-            "flash_attention": torch.nn.attention.SDPBackend.FLASH_ATTENTION,
-            "efficient_attention": torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION,
-            "cudnn_attention": torch.nn.attention.SDPBackend.CUDNN_ATTENTION,
-        }
-        return backend_map.get(backend_name.lower())
-    except AttributeError:
-        # torch.nn.attention.SDPBackend not available in older torch versions
-        return None
-
-
 def flush_memory():
     """Flush GPU memory and run garbage collection."""
     gc.collect()
@@ -223,12 +204,7 @@ def run_benchmark(
         self, model_id: str, config: BenchmarkConfig, num_tokens_to_profile: int = 0
     ) -> dict[str, Any] | None:
         """Run a single benchmark with the given model ID and config."""
-        sdpa_ctx = nullcontext()
-        if config.attn_implementation == "sdpa":
-            sdpa_backend = get_sdpa_backend(config.sdpa_backend)
-            sdpa_ctx = torch.nn.attention.sdpa_kernel(sdpa_backend)
-
-        with sdpa_ctx, torch.no_grad():
+        with torch.no_grad():
             self.logger.info(f"Running benchmark scenario: {config.name}")
 
             # Quick validation: try one measurement first to see if this scenario works
@@ -243,14 +219,14 @@ def run_benchmark(
 
             # Warmup runs
             self.logger.info(f"Warming up with {config.warmup_iterations} iterations...")
-            for _ in trange(config.warmup_iterations):
+            for _ in trange(config.warmup_iterations, desc="Warmup"):
                 _ = generate_fn(max_new_tokens=config.num_tokens_to_generate)
             self.logger.info("Warmup over.")
 
             # Measurement runs
             result = BenchmarkResult()
             self.logger.info(f"Benchmarking with {config.measurement_iterations} iterations.")
-            for _ in trange(config.measurement_iterations):
+            for _ in trange(config.measurement_iterations, desc="Benchmarking"):
                 e2e_latency, token_generation_times, shape_and_decoded_output, gpu_metrics = generate_fn(
                     max_new_tokens=config.num_tokens_to_generate,
                     gpu_monitor=(GPUMonitor(logger=self.logger) if config.gpu_monitoring else None),
@@ -281,49 +257,32 @@ def time_generate_batch(
     ) -> tuple[float, list[float], str, GPURawMetrics | None]:
         if gpu_monitor is not None:
             gpu_monitor.start()
-        config = GenerationConfig(
-            max_new_tokens=max_new_tokens,
-            eos_token_id=self.tokenizer.eos_token_id,
-            pad_token_id=self.tokenizer.pad_token_id,
-            do_sample=True,
-        )
-        manager = self.model.init_continuous_batching(config)
-        manager.start()
-        try:
-            first_req_results = []
-            timestamps = []
-            wall_time_0 = time.perf_counter()
-            inputs = self.inputs["input_ids"].tolist()
+        # Prepare inputs
+        inputs = self.inputs["input_ids"].tolist()
+        timestamps = []
+        last_result_generated_tokens = None
+        wall_time_0 = time.perf_counter()
+        # We disable prefix sharing because all prompts are the same
+        with self.model.continuous_batching_context_manager(allow_prefix_sharing=False) as manager:
             manager.add_requests(inputs, max_new_tokens=max_new_tokens, streaming=True)
-            first_req_id = None
-            num_requests = len(inputs)
-            finished_requests = 0
-            while finished_requests < num_requests:
+            unfinished_requests = len(inputs)
+            while unfinished_requests > 0:
                 # NOTE: I don't like having the extra if stmt here, but hopefully won't degrade perf too much
                 result = manager.get_result()
-                if result:
-                    timestamps.append(time.perf_counter() - wall_time_0)
+                if result is not None:
+                    timestamps.append(time.perf_counter() - wall_time_0)  # FIXME: the timestamps are wrong
                     if result.is_finished():
-                        finished_requests += 1
-                    if first_req_id is None:
-                        first_req_id = result.request_id
-                    if result.request_id == first_req_id:
-                        first_req_results.append(result)
-                else:
-                    if not manager.is_running():
-                        raise RuntimeError("Generation thread exited unexpectedly")
-            wall_time_1 = time.perf_counter()
-            gpu_metrics = gpu_monitor.stop_and_collect() if gpu_monitor is not None else None
-            decoded_output = self.tokenizer.decode(
-                [res.generated_tokens[0] for res in first_req_results], skip_special_tokens=True
-            )
-            shape_and_decoded_output = f"{(1, len(first_req_results))} | {decoded_output}"
-            e2e_latency = wall_time_1 - wall_time_0
-            return e2e_latency, timestamps, shape_and_decoded_output, gpu_metrics
-        except Exception as e:
-            raise e
-        finally:
-            manager.stop()
+                        last_result_generated_tokens = result.generated_tokens
+                        unfinished_requests -= 1
+                elif not manager.is_running():
+                    raise RuntimeError("Generation thread exited unexpectedly")
+        # Post-processing
+        wall_time_1 = time.perf_counter()
+        e2e_latency = wall_time_1 - wall_time_0
+        gpu_metrics = gpu_monitor.stop_and_collect() if gpu_monitor is not None else None
+        decoded_output = self.tokenizer.decode(last_result_generated_tokens, skip_special_tokens=True)
+        shape_and_decoded_output = f"{(1, len(last_result_generated_tokens))} | {decoded_output}"
+        return e2e_latency, timestamps, shape_and_decoded_output, gpu_metrics
 
     def time_generate(
         self,

filename: benchmark_v2/framework/data_classes.py
patch: @@ -89,10 +89,14 @@ class BenchmarkResult:
 
     def __init__(self) -> None:
         self.e2e_latency = []
-        self.token_generation_times = []  # time at which each token was generated (relative to start of the generation)
+        self.time_to_first_token = []
+        self.inter_token_latency = []
         self.shape_and_decoded_outputs = []
         self.gpu_metrics = []
 
+    def compute_itl(self, token_generation_times: list[float]) -> list[float]:
+        return (token_generation_times[-1] - token_generation_times[0]) / len(token_generation_times)
+
     def accumulate(
         self,
         e2e_latency: float,
@@ -101,7 +105,9 @@ def accumulate(
         gpu_metrics: GPURawMetrics | None,
     ) -> None:
         self.e2e_latency.append(e2e_latency)
-        self.token_generation_times.append(token_generation_times)
+        self.time_to_first_token.append(token_generation_times[0])
+        # inter-token latency is already an average in itself
+        self.inter_token_latency.append(self.compute_itl(token_generation_times))
         self.shape_and_decoded_outputs.append(shape_and_decoded_output)
         self.gpu_metrics.append(gpu_metrics)
 
@@ -113,7 +119,8 @@ def to_dict(self) -> dict[str, None | int | float]:
             gpu_metrics = [gm.to_dict() for gm in self.gpu_metrics]
         return {
             "e2e_latency": self.e2e_latency,
-            "token_generation_times": self.token_generation_times,
+            "time_to_first_token": self.time_to_first_token,
+            "inter_token_latency": self.inter_token_latency,
             "shape_and_decoded_outputs": self.shape_and_decoded_outputs,
             "gpu_metrics": gpu_metrics,
         }
@@ -130,29 +137,25 @@ def from_dict(cls, data: dict[str, None | int | float]) -> "BenchmarkResult":
         for i in range(len(data["e2e_latency"])):
             new_instance.accumulate(
                 e2e_latency=data["e2e_latency"][i],
-                token_generation_times=data["token_generation_times"][i],
+                time_to_first_token=data["time_to_first_token"][i],
+                inter_token_latency=data["inter_token_latency"][i],
                 shape_and_decoded_output=data["shape_and_decoded_outputs"][i],
                 gpu_metrics=gpu_metrics[i],
             )
         return new_instance
 
-    def get_measured_ttft(self) -> list[float]:
-        return [dt[0] for dt in self.token_generation_times if len(dt) > 0]
-
-    def get_measured_itl(self) -> list[float]:
-        return [(dt[-1] - dt[0]) / (len(dt) - 1) for dt in self.token_generation_times if len(dt) > 1]
-
     def get_throughput(self, total_generated_tokens: int) -> list[float]:
         return [total_generated_tokens / e2e_latency for e2e_latency in self.e2e_latency]
 
     def pprint(self, batch_size: int = 0, num_generated_tokens: int = 0, tabs: int = 0) -> None:
         measurements = {
             "E2E Latency": add_unit_to_duration(compute_basic_statistics(self.e2e_latency)),
-            "Time to First Token": add_unit_to_duration(compute_basic_statistics(self.get_measured_ttft())),
+            "Time to First Token": add_unit_to_duration(compute_basic_statistics(self.time_to_first_token)),
         }
-        itl_values = self.get_measured_itl()
-        if len(itl_values) > 0:
-            measurements["Inter-Token Latency"] = add_unit_to_duration(compute_basic_statistics(itl_values))
+        if len(self.inter_token_latency) > 0:
+            measurements["Inter-Token Latency"] = add_unit_to_duration(
+                compute_basic_statistics(self.inter_token_latency)
+            )
         if batch_size > 0:
             throughput_stats = compute_basic_statistics(self.get_throughput(batch_size * num_generated_tokens))
             measurements["Throughput"] = {key: f"{value:.2f}tok/s" for key, value in throughput_stats.items()}

filename: benchmark_v2/run_benchmarks.py
patch: @@ -19,19 +19,20 @@
 """
 
 import argparse
+import json
 import logging
 import sys
 import uuid
 
-from framework.benchmark_config import adapt_configs, get_config_by_level
+from framework.benchmark_config import BenchmarkConfig, adapt_configs, get_config_by_level
 from framework.benchmark_runner import BenchmarkRunner
 
 
 if __name__ == "__main__":
     # Parse arguments
     parser = argparse.ArgumentParser()
     parser.add_argument("--output-dir", type=str, default=None, help="Output dir for benchmark results")
-    parser.add_argument("--log-level", type=str, choices=["DEBUG", "INFO", "WARNING", "ERROR"], default="INFO")
+    parser.add_argument("--log-level", type=str, choices=["DEBUG", "INFO", "WARNING", "ERROR"], default="WARNING")
     parser.add_argument("--model-id", type=str, help="Specific model ID to benchmark (if supported by benchmarks)")
     parser.add_argument("--warmup", "-w", type=int, default=3, help="Number of warmup iterations")
     parser.add_argument("--iterations", "-i", type=int, default=10, help="Number of measurement iterations")
@@ -48,6 +49,7 @@
         " each attn implementation an option, 3: cross-generate all combinations of configs, 4: cross-generate all"
         " combinations of configs w/ all compile modes",
     )
+    parser.add_argument("--config-file", type=str, help="Path to a config file stored as a json or jsonl format")
     parser.add_argument("--num-tokens-to-profile", "-p", type=int, default=0, help="Number of tokens to profile")
 
     parser.add_argument("--branch-name", type=str, help="Git branch name")
@@ -90,8 +92,21 @@
     if any(n <= 1 for n in args.num_tokens_to_generate):
         raise ValueError("--num_tokens_to_generate arguments should be larger than 1")
 
-    # Get the configs for the given coverage level
-    configs = get_config_by_level(args.level)
+    # If a config file is provided, read it and use the configs therein. They will still be adapted to the given arguments.
+    if args.config_file is not None:
+        if args.config_file.endswith(".json"):
+            with open(args.config_file, "r") as f:
+                config_as_dicts = [json.load(f)]
+        elif args.config_file.endswith(".jsonl"):
+            with open(args.config_file, "r") as f:
+                config_as_dicts = [json.loads(line) for line in f if line.startswith("{")]
+        else:
+            raise ValueError(f"Unsupported config file format: {args.config_file}")
+        configs = [BenchmarkConfig.from_dict(config) for config in config_as_dicts]
+    else:
+        # Otherwise, get the configs for the given coverage level
+        configs = get_config_by_level(args.level)
+
     # Adapt the configs to the given arguments
     configs = adapt_configs(
         configs,

filename: src/transformers/generation/continuous_batching/continuous_api.py
patch: @@ -826,8 +826,6 @@ def stop(self, block: bool = True, timeout: float | None = None) -> None:
                 logger.warning(
                     f"\nPrefix sharing was on. Total prefix length: {self.batch_processor.cache._total_prefix_length}"
                 )
-            else:
-                logger.warning("\nPrefix sharing was off.")
 
         if self._generation_thread is None:
             logger.warning("Manager not started.")