add attn bench #24
add attn bench #24
Conversation
Summary of ChangesHello @yiliu30, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a comprehensive benchmarking suite for attention mechanisms, specifically designed for XPU devices. It provides a structured way to evaluate the performance of various attention implementations, such as SDPA, FlexAttention, and Eager attention, by measuring both forward and backward pass execution times. The inclusion of a dedicated XPU profiler wrapper enhances the ability to gather detailed performance metrics for XPU-accelerated operations. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds a new benchmark for attention mechanisms. My review found a few issues. The main benchmark script contains a hardcoded model path, which harms portability, and has several blocks of commented-out code that should either be enabled via configuration or removed. The accompanying profiler utility script, which appears to be adapted from another project, contains a critical bug where it calls a non-existent logger method, which will cause a runtime error. Finally, the shell script to run the benchmark is missing a standard trailing newline. I've provided specific suggestions to address these points.
| logger.info_once( | ||
| "GPU profiling will start " | ||
| f"{self._delay_iters} steps after start_profile." | ||
| ) |
There was a problem hiding this comment.
The loguru.logger object does not have a info_once method. This will raise an AttributeError at runtime. This seems to be a leftover from copying the code from vLLM, which uses a custom logger setup. Please replace info_once with the standard info method.
| logger.info_once( | |
| "GPU profiling will start " | |
| f"{self._delay_iters} steps after start_profile." | |
| ) | |
| logger.info( | |
| "GPU profiling will start " | |
| f"{self._delay_iters} steps after start_profile." | |
| ) |
| logger.info_once( | ||
| "GPU profiling will stop " | ||
| f"after {self._max_iters} worker steps, " | ||
| "or when stop_profile is received." | ||
| ) |
There was a problem hiding this comment.
The loguru.logger object does not have a info_once method. This will raise an AttributeError at runtime. Please replace it with the standard info method.
| logger.info_once( | |
| "GPU profiling will stop " | |
| f"after {self._max_iters} worker steps, " | |
| "or when stop_profile is received." | |
| ) | |
| logger.info( | |
| "GPU profiling will stop " | |
| f"after {self._max_iters} worker steps, " | |
| "or when stop_profile is received." | |
| ) |
|
|
||
| def shutdown(self) -> None: | ||
| """Ensure profiler is stopped when shutting down.""" | ||
| logger.info_once("Shutting down profiler") |
There was a problem hiding this comment.
The loguru.logger object does not have a info_once method. This will raise an AttributeError at runtime. Please replace it with the standard info method.
| logger.info_once("Shutting down profiler") | |
| logger.info("Shutting down profiler") |
| from loguru import logger | ||
|
|
||
| # --- 1. Setup Configuration --- | ||
| model_name = "/data/yiliu/models/Qwen/Qwen3-8B" |
There was a problem hiding this comment.
The model path is hardcoded. This makes the script less portable and harder to use for other models or in different environments. It's better to pass this as a command-line argument, for example using Python's argparse module.
| total_flops_per_step = fwd_flops if forward_only else (fwd_flops * 3) | ||
|
|
||
| tflops = (total_flops_per_step / (avg_time_per_step / 1000)) / 1e12 | ||
| # logger.info(f" > Approx TFLOPS: {tflops:.2f}") |
There was a problem hiding this comment.
The TFLOPS calculation is performed, but the result is never used because this logging statement is commented out. This makes the calculation on the preceding lines dead code. Please uncomment this line to display the TFLOPS.
| # logger.info(f" > Approx TFLOPS: {tflops:.2f}") | |
| logger.info(f" > Approx TFLOPS: {tflops:.2f}") |
| # logger.info("=== MODE: FORWARD ONLY ===") | ||
| # for impl in implementations: | ||
| # torch.xpu.empty_cache() | ||
| # benchmark_implementation(impl, forward_only=True) |
There was a problem hiding this comment.
This block for the 'Forward Only' benchmark is commented out. If this mode is useful, it should be enabled (e.g., via a command-line argument) rather than being commented out. If it's no longer needed, it should be removed. I'm suggesting to uncomment it for now.
| # logger.info("=== MODE: FORWARD ONLY ===") | |
| # for impl in implementations: | |
| # torch.xpu.empty_cache() | |
| # benchmark_implementation(impl, forward_only=True) | |
| logger.info("=== MODE: FORWARD ONLY ===") | |
| for impl in implementations: | |
| torch.xpu.empty_cache() | |
| benchmark_implementation(impl, forward_only=True) |
| @@ -0,0 +1 @@ | |||
| ZE_AFFINITY_MASK=3 python bench_attn.py No newline at end of file | |||
There was a problem hiding this comment.
Shell script files should end with a newline character for POSIX compatibility. Some tools might not process the last line correctly without it.
| ZE_AFFINITY_MASK=3 python bench_attn.py | |
| ZE_AFFINITY_MASK=3 python bench_attn.py | |
Signed-off-by: yiliu30 <[email protected]>
Signed-off-by: yiliu30 [email protected]