Adding FP8 benchmark on attention and matmul testing

[bkryu]

📌 Description

Current PR extends benchmarking script in flashinfer_benchmark.py by adding:

• MLA backend testing
• Attention FP8 Benchmarking Introduction for prefill, decode, and MLA.
• bmm_fp8 and mm_fp4 testing.
• Addition of a few backends where applicable (e.g. trtllm-gen for paged prefill)
• Fix broken cuDNN prefill & decode benchmarking due to migration from cubins to full cuDNN integration
• General minor benchmark code refactoring to reduce code redundancy.

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

[gemini-code-assist]

Summary of Changes

Hello @bkryu, I'm Gemini Code Assist¹! 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 significantly enhances our benchmarking suite by integrating support for FP8 operations. It introduces new benchmarks for Multi-Layer Attention and Batch Matrix Multiplication, alongside updating existing attention benchmarks to handle FP8 data types. These changes provide a robust framework for evaluating the performance of low-precision computations.

Highlights

• FP8 Benchmarking Introduction: I've added comprehensive benchmarking support for FP8 (8-bit floating point) operations across both attention and matrix multiplication routines. This allows for performance evaluation of low-precision computations.
• New Attention Benchmark (MLA): A new benchmark, BatchMLAPagedAttentionWrapper, has been introduced to specifically test Multi-Layer Attention (MLA) with paged KV cache, including support for FP8 data types.
• New GEMM Benchmark (FP8 BMM): I've implemented a new benchmark, bmm_fp8, for evaluating the performance of FP8 Batch Matrix Multiplication (BMM) using various backends like cuDNN and cuBLAS.
• FP8 Data Type Support & Backend Compatibility: Existing attention benchmarks (BatchDecodeWithPagedKVCacheWrapper, BatchPrefillWithPagedKVCacheWrapper, BatchPrefillWithRaggedKVCacheWrapper) have been updated to support FP8 query and KV cache data types, with logic to skip incompatible backends (e.g., FA2_TC, cuDNN for FP8).
• Benchmark Infrastructure Refinements: The benchmark routines have been refactored to improve code organization, including consolidating wrapper initialization into dictionaries and introducing a shared print_perf_metrics utility function for consistent output.

Using Gemini Code Assist

The 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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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 or fill out our survey to provide feedback.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request adds new FP8 benchmarks for attention and matrix multiplication routines. The changes are functionally correct and expand the test coverage nicely. My main feedback focuses on improving maintainability by addressing code duplication. There are several opportunities to refactor repeated logic (like routine lists, data type parsing, and backend capability checks) into shared constants and helper functions. This will make the benchmark suite easier to read, maintain, and extend in the future.

[yzh119]

Is it because we directly call the trtllm_gen decode functions instead of going through the wrapper?

[bkryu]

Yes. As stated in the documentation for BatchMLAPagedAttentionWrapper, supported backends are fa2 & fa3 (+ auto).

trtllm-gen directly needs to go to through flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla. However, I did not get a chance to have it supported for flashinfer_benchmark.py yet due to my bandwidth limitations. I initially missed it due to trtllm_batch_decode_with_kv_cache_mla being omitted in our official documentation (and we may want to reconsider moving it from flashinfer.decode to flashinfer.mla )

[yzh119]

no problem, thanks for the explaination!

[yzh119]

LGTM, thank you @bkryu !