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.

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

[TUTORIAL] non-causal mode in fused-attention backward #5241

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

arthurfeeney
Copy link

Description

This makes some simple additions to add support for non-causal attention in 06-fused-attention.py. The forward pass already uses the causal parameter, so this just changes the backward pass.

I think having the non-causal mode supported makes the code a little easier to understand. I.e., I did not find it obvious why the backward pass always masked the diagonal blocks until I realized it was just always computing the causal version.

New contributor declaration

  • I am not making a trivial change, such as fixing a typo in a comment.

  • I have written a PR description following these
    rules.

  • I have run pre-commit run --from-ref origin/main --to-ref HEAD.

  • Select one of the following.

    • I have added tests.
      • /test for lit tests
      • /unittest for C++ tests
      • /python/test for end-to-end tests
    • This PR does not need a test because The script includes a simple test, so I set it to run the non-causal mode as well.
  • Select one of the following.

    • I have not added any lit tests.
    • The lit tests I have added follow these best practices,
      including the "tests should be minimal" section. (Usually running Python code
      and using the instructions it generates is not minimal.)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant