Refactor benchmarks for Flash Attention Prefill #447
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
muhammad-tanvir-1211
force-pushed
the
flash_prefill_separate_out_benchmarks
branch
2 times, most recently
from
c863845 to
e7ec638
Compare
* Add comment on final output type conversion
muhammad-tanvir-1211
force-pushed
the
flash_prefill_separate_out_benchmarks
branch
from
e7ec638 to
6a7ad57
Compare
aacostadiaz
added a commit
that referenced
this pull request
Jun 28, 2025
This PR separates the output type and accumulator type for Flash Attention Prefill. Combinations supported are: * bf16 inputs, fp32 accumulator, bf16 | fp32 output * fp16 inputs, fp32 accumulator, fp16 | fp32 output * fp8 inputs, fp32 accumulator, fp8 | fp32 output Tests added in: #446 Benchmarks added in: #447 --------- Co-authored-by: Alejandro Acosta <[email protected]>
t4c1
reviewed
Jul 2, 2025
applications/flash_attention_v2/collective/xe_flash_attn_prefill_epilogue.hpp
Show resolved
Hide resolved
| @@ -77,12 +77,17 @@ struct FMHAOptions { | |||
| cmd.get_cmd_line_argument("num_heads_kv", num_heads_kv, num_heads_q); | |||
| cmd.get_cmd_line_argument("seq_len_qo", seq_len_qo, 512); | |||
| cmd.get_cmd_line_argument("seq_len_kv", seq_len_kv, seq_len_qo); | |||
| cmd.get_cmd_line_argument("head_size_vo", head_size_vo, 128); | |||
| cmd.get_cmd_line_argument("head_size_vo", head_size_vo, HEAD_DIM); | |||
There was a problem hiding this comment.
Why even keep this as an option if it must be set to a certain value?
There was a problem hiding this comment.
This helps throw an error on line 88 for the case when the head_size_vo received by a benchmark does not match the head_size_vo this benchmark was built for. For example, if we pass head_size_vo=96 to the cutlass_benchmark_flash_attention_prefill_h64_xe exe file, it should not run the benchmark and throw an error instead.
There was a problem hiding this comment.
What I am trying to point out is that given we have the head size in the name of the benchmark, we should not need an argument for head size at all. Without the argument there is no chance of passing wrong argument value.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR adds benchmarks for all the different data types supported with Flash Attention Prefill. It is a continuation of PR #443