[JAX] Fused attention unit tests fixes and refinements #1352
Conversation
This file contains 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
zlsh80826
changed the title
|
/te-ci jax L1 |
zlsh80826
force-pushed
the
rewang/fa-refactor
branch
from
42c31ab to
28e4ded
Compare
|
/te-ci jax L1 |
zlsh80826
force-pushed
the
rewang/fa-refactor
branch
from
28e4ded to
76780a3
Compare
|
/te-ci jax L1 |
1 similar comment
|
/te-ci jax L1 |
zlsh80826
force-pushed
the
rewang/fa-refactor
branch
from
76780a3 to
7c76d81
Compare
|
/te-ci jax L1 |
zlsh80826
force-pushed
the
rewang/fa-refactor
branch
from
7c76d81 to
fb2c73b
Compare
Signed-off-by: Reese Wang <[email protected]>
Signed-off-by: Reese Wang <[email protected]>
Signed-off-by: Reese Wang <[email protected]>
Signed-off-by: Reese Wang <[email protected]>
Signed-off-by: Reese Wang <[email protected]>
Signed-off-by: Reese Wang <[email protected]>
Signed-off-by: Reese Wang <[email protected]>
Signed-off-by: Reese Wang <[email protected]>
Signed-off-by: Reese Wang <[email protected]>
Signed-off-by: Reese Wang <[email protected]>
zlsh80826
force-pushed
the
rewang/fa-refactor
branch
from
fb2c73b to
74f0603
Compare
|
/te-ci jax L1 |
mgoldfarb-nvidia
approved these changes
Dec 16, 2024
| @@ -658,16 +668,6 @@ def check_dqkv(primitive, reference, pad): | |||
| ) | |||
|
|
|||
|
|
|||
| @pytest.mark.parametrize( | |||
There was a problem hiding this comment.
Good catch - the number of skipped tests was getting very long
phu0ngng
approved these changes
Dec 16, 2024
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
Refinments for the JAX test_fused_attn and fix a THD cross attention + causal bug.
Type of change
Changes
segment_possegment_padin the unit tests, and encoded the padding informations to thesegment_ids. (0s means paddings)test_praxis_layersspecific, which allows test_fused_attn.py to run with the non-deterministic (faster and widely used) kernels."attn_bias_type, bias_shape"for forward and backward to avoid lots of skipped tests.cachedecorator for fused attn deterministic, which will cause errors on the unit tests.Checklist: