Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Dynamo/AOTAutograd traceable flash attention #8654

Merged
merged 4 commits into from
Feb 1, 2025
Merged

Dynamo/AOTAutograd traceable flash attention #8654

merged 4 commits into from
Feb 1, 2025
+707 −285

Conversation

zpcore
Copy link
Collaborator

@zpcore zpcore commented Jan 30, 2025 •
edited
Loading

Resolves #8633

@zpcore zpcore force-pushed the piz/autograde_trace branch from f82f373 to a8c8f47 Compare January 31, 2025 10:31
@zpcore zpcore changed the title backward with spmd issue Dynamo/AOTAutograd traceable flash attention Jan 31, 2025
@zpcore zpcore requested a review from tengyifei January 31, 2025 10:38
tengyifei
tengyifei approved these changes Jan 31, 2025
View reviewed changes
Copy link
Collaborator

@tengyifei tengyifei left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is really great

test/test_pallas.py Show resolved Hide resolved
torch_xla/experimental/custom_kernel.py Outdated Show resolved Hide resolved
torch_xla/experimental/custom_kernel.py Show resolved Hide resolved
torch_xla/experimental/custom_kernel.py Outdated Show resolved Hide resolved
torch_xla/experimental/custom_kernel.py Show resolved Hide resolved
torch_xla/experimental/custom_kernel.py Show resolved Hide resolved
test/test_pallas_spmd.py Outdated Show resolved Hide resolved
test/test_pallas_spmd.py Show resolved Hide resolved
test/test_pallas.py Outdated Show resolved Hide resolved
test/test_pallas.py Outdated Show resolved Hide resolved
@zpcore zpcore merged commit 9ae017e into master Feb 1, 2025
12 checks passed
@zpcore zpcore deleted the piz/autograde_trace branch February 1, 2025 04:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@tengyifei tengyifei tengyifei approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

Make flash_attention Dynamo/AOTAutograd traceable
2 participants
@zpcore @tengyifei