Add torch.compile support to SageAttention

[guilhermeleobas]

I'm not 100% sure if the torch compile tests are representative of a common SageAttention usage. I wrote them based on the benchmark files.

Also, is there any lint rule that I should apply to the files?

Edit: I just saw there's a beta version of SageAttention 3 on HuggingFace. Is the code available? If so, I can also work on add torch.compile support on it.

[guilhermeleobas]

cc @StrongerXi

[StrongerXi]

@jt-zhang @jason-huang03 would it be possible to merge this patch? It's pretty harmless, most of the code is adding fake impls for the sage attention ops (you can think of them as registering shape functions for torch.compile).

It actually fixes some accuracy issue which showed up in sage + compile: comfyanonymous/ComfyUI#8689 (comment).

Also, diffusers is adding Sage Attention as a backend: huggingface/diffusers#12439, so merging this patch would benefit many users who often use sage + compile in diffusers and comfyui.

[jt-zhang]

Thank you for your pr. Please @whx1003 help to check and merge this pr.

[sayakpaul]

Hey folks!

I am one of the maintainers of diffusers and can vouch for the merit of this PR as it benefits the torch.compile users quite a bit.

This would be great to have merged. Also cc: @MekkCyber

[whx1003]

@guilhermeleobas Thanks for the PR! The code looks good to me — I ran tests on a RTX5090 and everything worked as expected.

Could you please keep only the changes under the sageattention/ directory in this PR? The other files seem unrelated.

[sayakpaul]

@whx1003 tests/test_torch_compile.py could be beneficial for compilation tests. WDYT?

[whx1003]

@sayakpaul I agree that the tests could be useful, but we don’t maintain test files in this repo at the moment.

I’d prefer to leave them out of this PR and maybe revisit adding them later.

[sayakpaul]

No worries! @guilhermeleobas maybe we can have the test file as a gist under your profile and mention a link.

[guilhermeleobas]

Thanks for the feedback folks. @whx1003 I've removed the test file and unrelated changes.

[whx1003]

Thanks!

[woct0rdho]

nitpick: If I understand correctly, in the fake impl:

SageAttention/sageattention/sm80_compile.py

Line 121 in 15c0e22

lse = torch.empty((batch_size, num_qo_heads, qo_len), dtype=torch.float32, device="cuda")

it's better to set lse's device to query.device rather than "cuda" (the 0-th device), so it's consistent with the C code:

SageAttention/csrc/qattn/qk_int_sv_f16_cuda_sm80.cu

Line 779 in 15c0e22

lse = torch::empty({batch_size, num_qo_heads, qo_len}, query.options().dtype(torch::kFloat32));

This is also true for sm89 and sm90. Maybe this helps solve the issue of set_device:

SageAttention/sageattention/core.py

Lines 250 to 256 in 15c0e22

	# FIXME(DefTruth): make sage attention work compatible with distributed
	# env, for example, xDiT which launch by torchrun. Without this workaround,
	# sage attention will run into illegal memory access error after first
	# inference step in distributed env for multi gpus inference. This small
	# workaround also make sage attention work compatible with torch.compile
	# through non-fullgraph compile mode.
	torch.cuda.set_device(v.device)

but I haven't tested.

[woct0rdho]

Another thing noteworthy, although I haven't tested, is that we may need to set mutates_args={"output"}

If I understand correctly, this should be absolutely needed, but I couldn't yet find an example that gives correct output with this and wrong output without this.

In #74 (comment) , he succeeded by specifying q, k, v as mutable parameters, but I think that makes no sense...