Skip to content

[WIP][aiter][flash_attention] update aiter and add attn #612

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

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

[WIP][aiter][flash_attention] update aiter and add attn #612

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e074771
update aiter
xuzhao9 Nov 1, 2025
1324d10
fix lint
xuzhao9 Nov 1, 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
2 changes: 1 addition & 1 deletion submodules/aiter
Open in desktop
Submodule aiter updated 2771 files
17 changes: 17 additions & 0 deletions tritonbench/operators/flash_attention/operator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,9 @@

from .test_fmha_utils import make_packed_qkv

with try_import("HAS_AITER"):
from aiter.ops.triton.mha import flash_attn_func as aiter_flash_attn_func

HAS_CUDA_124 = (
torch.cuda.is_available() and torch.version.cuda and torch.version.cuda >= "12.4"
)
Expand Down Expand Up @@ -362,6 +365,20 @@ def sdpa_flash_attention(q, k, v):
v,
)

@register_benchmark(enabled=is_hip() and HAS_AITER)
def aiter(self, q, k, v):
def _inner():
return aiter_flash_attn_func(
q,
k,
v,
softmax_scale=self.sm_scale,
causal=self.causal,
deterministic=self.deterministic,
)

return _inner

if IS_B200:
# Only enable calling this benchmark directly.
@register_benchmark(enabled=False)
Expand Down