JAX FA Benchmarking Script #351
Open
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
Micky774
requested review from
ipanfilo,
wangye805 and
wenchenvincent
as code owners
ipanfilo
reviewed
Oct 31, 2025
transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp
Outdated
Show resolved
Hide resolved
ipanfilo
reviewed
Oct 31, 2025
Collaborator
ipanfilo
left a comment
ipanfilo
left a comment
There was a problem hiding this comment.
Why ck_fused_attn_bwd needs modification if the PR is for fwd-pass only?
Micky774
changed the title
Contributor
Author
|
Note that I have added the BWD pass implementation as well to this PR. |
Contributor
Author
|
Pinging @wangye805 @wenchenvincent in case either of you are interested in reviewing this PR as well, thanks! |
benchmarks/attention/README.md
Outdated
| ## JAX Fused-Attention Benchmarking | ||
| The benchmarking process is split into two stages: *generating* the timing data, and *visualizing* the timing data. The following steps assume you are located in `TransformerEngine/benchmarks/attention` (i.e. where this README is located). First, ensure that you install requirements via `pip install -r requirements.txt`. | ||
|
|
||
| Note: Only forward timings are supported at this point. |
| from transformer_engine.jax import fp8_autocast | ||
|
|
||
| # Needed in order to dump timings properly | ||
| os.environ["XLA_FLAGS"]="--xla_gpu_graph_level=0" |
Collaborator
There was a problem hiding this comment.
Is this because you used dumping time function in ck fused attention?
ipanfilo
reviewed
Nov 14, 2025
ipanfilo
approved these changes
Dec 1, 2025
| attn_bias_type, bias_shape = bias_config | ||
| window_size = None | ||
| if swa: | ||
| window_size = (s_kv // 10, 0) |
Collaborator
There was a problem hiding this comment.
Why do this for SWA?
Contributor
Author
There was a problem hiding this comment.
This was taken from our JAX FA testing.
Collaborator
|
Have you compared the kernel time measured from CK FA API vs from rocprof? |
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.
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: