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

Restore the support of causal=True and seqlen_q != seqlen_k #55

Open
wants to merge 14 commits into
base: main
Choose a base branch
from
Open

Restore the support of causal=True and seqlen_q != seqlen_k #55

wants to merge 14 commits into from

Conversation

xinyazhang
Copy link
Collaborator

@xinyazhang xinyazhang commented Nov 14, 2024 •
edited
Loading

Major changes

  1. Fix the Triton kernel to support causal=True and seqlen_q != seqlen_k
    • This matches efficient attention behavior (align to top left corner)
  2. Re-Enable the related tests
  3. Re-Enable the related tuning (also fixes OOM on long seqlen_q/k)
  4. Add the missing entries to tuning database for causal=True and seqlen_q != seqlen_k

Known Problems

No tuning entries for MI300X/Navi31 since we are going to re-run the tuning script for all cases no later than next week.

@xinyazhang
Fix causal mask for fwd kernel.
279e0d5 
Need top left alignment.
@xinyazhang
Re-Enable casual=True when seqlen_q != seqlen_k
f9cdb64
@xinyazhang
Call torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp when avai…
8c68bc4 
…lable
@xinyazhang
Bare minimal tests for casual=True when seqlen_q != seqlen_k
0bcd5a1
@xinyazhang
wip
8555423
@xinyazhang
Another fix of causal mask for fwd kernel.
543649e 
NaN masking has some problem.
@xinyazhang
Remove skip code in bwd kernel for casual=True and seqlen_q != seqlen_k
5beb80a
@xinyazhang
Revert "wip"
538b696 
This reverts commit 2079b3b.
@xinyazhang
Fix the algorithm to compute leading/full/trailing blocks for dq
12d2561
@xinyazhang
tune_flash: Unskip casual=True and seqlen_q != seqlen_k
1a40610
@xinyazhang
tune_flash: add --causal_non_square_only and fix OOM
95b424a 
Need to reduce batch size and number of heads for long sequences.
@xinyazhang
Require full LUT tensor for causal=True and bias=0
654569d
@xinyazhang
Add missing causal=True entries to the tuning database for MI200 GPU
9fbfe26 
CAVEAT: MI300X/Navi31's tuning entries are still incomplete.
@xinyazhang
Defer raise MissingLutEntry so we can have some partial output to par…
6ca821d 
…se the result.
@xinyazhang xinyazhang marked this pull request as ready for review November 14, 2024 05:17
@xinyazhang
Copy link
Collaborator Author

xinyazhang commented Nov 14, 2024 •
edited
Loading

Most tests passed (105 failed vs 36759 passed).
The failed tests are due to higher numerical errors on MI200 hardware (and I didn't bother to fine tune the tolerance for MI200 right now).

image

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@pruthvistony pruthvistony pruthvistony approved these changes

@jithunnair-amd jithunnair-amd Awaiting requested review from jithunnair-amd

@vgokhale vgokhale Awaiting requested review from vgokhale

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

Successfully merging this pull request may close these issues.
2 participants
@xinyazhang @pruthvistony