Skip to content

Add torch.compile support to flash attention 3 #1769

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

Merged
v0i0 merged 18 commits into from
Dec 4, 2025
Merged

Add torch.compile support to flash attention 3 #1769

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
c7697bb
Add torch.compile support to flash attention 3
guilhermeleobas Jul 16, 2025
e1944ba
Don't return mutated variables in mha_bwd
guilhermeleobas Jul 24, 2025
a760ca3
Change fake_check flag to be opt-in; Remove build.sh and remove if-el…
guilhermeleobas Jul 25, 2025
24cc2b2
Remove print statements and update exception message
guilhermeleobas Jul 30, 2025
5e114d5
Fix flash_attn_backward_fake
guilhermeleobas Aug 6, 2025
734bc43
Add `safe_aot_autograd_check`
guilhermeleobas Aug 7, 2025
fde4bc0
Update namespace to flash_attn_3
guilhermeleobas Aug 19, 2025
ab79ae2
Add `flash_attn_forward.register_autograd`
guilhermeleobas Aug 22, 2025
6250fbe
Fix bug in `flash_attn_backward_fake`
guilhermeleobas Aug 22, 2025
1e3539e
Add support and tests for torch.export and aoti_compile_and_package
guilhermeleobas Sep 2, 2025
f174bd6
format code
guilhermeleobas Sep 3, 2025
6fe1c8c
update flash_api_stable.cpp
guilhermeleobas Sep 19, 2025
b555ac7
Fix flash_api_stable.cpp build
guilhermeleobas Oct 13, 2025
0aa4fa1
Only run schema_check if dtype is not float8_e4m3fn
guilhermeleobas Oct 13, 2025
47d7137
Correctly compute kBlockM for sm88/86/80
guilhermeleobas Oct 13, 2025
49fb775
Fix bug in boxed_mha_bwd
guilhermeleobas Oct 13, 2025
65dd580
don't run autograd_check when num_splits > 0
guilhermeleobas Nov 12, 2025
3fcde4b
Raise TypeError if out is specified when compiling _flash_attn_forward
guilhermeleobas Nov 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,10 @@ var/

# IDE-related
.idea/
.vscode/

# Dev
venv

# compile-time generated file
flash_attn_config.py
6 changes: 3 additions & 3 deletions hopper/flash_api.cpp
View file
Open in desktop
Copy link
Copy Markdown
Contributor

@janeyx99 janeyx99 Sep 15, 2025

Choose a reason for hiding this comment

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

Now that #1791 is landed can you also update the stable file to keep them in sync?

Copy link
Copy Markdown
Contributor Author

@guilhermeleobas guilhermeleobas Sep 19, 2025

Choose a reason for hiding this comment

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

done!

Original file line number Diff line number Diff line change
Expand Up @@ -1264,7 +1264,7 @@ void run_mha_bwd(Flash_bwd_params &params, cudaStream_t stream) {
// h: num_heads
// h_k: num_heads_k
// d: head_size
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_bwd(
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_bwd(
at::Tensor dout, // (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q
at::Tensor q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
at::Tensor k, // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
Expand Down Expand Up @@ -1563,7 +1563,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tenso
softmax_d.zero_();
}

return { dq, dk, dv, softmax_d, softmax_lse_log2, dq_accum, dk_accum, dv_accum };
return { softmax_d, softmax_lse_log2, dq_accum, dk_accum, dv_accum };
}

std::tuple<at::Tensor, at::Tensor>
Expand Down Expand Up @@ -1727,7 +1727,7 @@ TORCH_LIBRARY(flash_attn_3, m) {
"int window_size_right = -1,"
"float softcap = 0.0,"
"bool deterministic = False,"
"int sm_margin = 0) -> (Tensor(dq!), Tensor(dk!), Tensor(dv!), Tensor, Tensor, Tensor, Tensor, Tensor)");
"int sm_margin = 0) -> (Tensor, Tensor, Tensor, Tensor, Tensor)");
m.def("fwd_combine("
"Tensor out_partial,"
"Tensor lse_partial,"
Expand Down
21 changes: 9 additions & 12 deletions hopper/flash_api_stable.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1335,7 +1335,7 @@ void run_mha_bwd(Flash_bwd_params &params, cudaStream_t stream) {
// h: num_heads
// h_k: num_heads_k
// d: head_size
std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor> mha_bwd(
std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor> mha_bwd(
Tensor dout, // (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q
Tensor q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
Tensor k, // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
Expand Down Expand Up @@ -1641,7 +1641,7 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor> mha_b
torch::stable::zero_(softmax_d);
}

return { dq, dk, dv, softmax_d, softmax_lse_log2, dq_accum, dk_accum, dv_accum };
return { softmax_d, softmax_lse_log2, dq_accum, dk_accum, dv_accum };
}

std::tuple<Tensor, Tensor>
Expand Down Expand Up @@ -1828,16 +1828,13 @@ void boxed_mha_bwd(
auto deterministic = to<bool>(stack[20]);
auto sm_margin = to<int64_t>(stack[21]);

auto [dq_, dk_, dv_, softmax_d, softmax_lse_log2, dq_accum, dk_accum, dv_accum] = mha_bwd(dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, max_seqlen_q, max_seqlen_k, softmax_scale, is_causal, window_size_left, window_size_right, softcap, deterministic, sm_margin);
auto [softmax_d, softmax_lse_log2, dq_accum, dk_accum, dv_accum] = mha_bwd(dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, max_seqlen_q, max_seqlen_k, softmax_scale, is_causal, window_size_left, window_size_right, softcap, deterministic, sm_margin);

stack[0] = from(dq_);
stack[1] = from(dk_);
stack[2] = from(dv_);
stack[3] = from(softmax_d);
stack[4] = from(softmax_lse_log2);
stack[5] = from(dq_accum);
stack[6] = from(dk_accum);
stack[7] = from(dv_accum);
stack[0] = from(softmax_d);
stack[1] = from(softmax_lse_log2);
stack[2] = from(dq_accum);
stack[3] = from(dk_accum);
stack[4] = from(dv_accum);
}

void boxed_mha_combine(
Expand Down Expand Up @@ -1949,7 +1946,7 @@ STABLE_TORCH_LIBRARY(flash_attn_3, m) {
"int window_size_right = -1,"
"float softcap = 0.0,"
"bool deterministic = False,"
"int sm_margin = 0) -> (Tensor(dq!), Tensor(dk!), Tensor(dv!), Tensor, Tensor, Tensor, Tensor, Tensor)");
"int sm_margin = 0) -> (Tensor, Tensor, Tensor, Tensor, Tensor)");
m.def("fwd_combine("
"Tensor out_partial,"
"Tensor lse_partial,"
Expand Down
Loading