Skip to content

Added Perf Model for aiter::flash_attn#242

Merged
ajassani merged 5 commits intomainfrom
spanmore/add_aiter_fa_perf_model
Jul 24, 2025
Merged

Added Perf Model for aiter::flash_attn#242
ajassani merged 5 commits intomainfrom
spanmore/add_aiter_fa_perf_model

Conversation

@spandoesai
Copy link
Copy Markdown
Collaborator

Added perf analysis support for the flash attention implementation in aiter for both forward and backward passes.

Changes made:

  • Define custom Perf Model classes for aiter::_flash_attn_fwd and aiter::_flash_attn_bwd perf_model.py
  • Link the ops to the appropriate Perf Model torch_op_mapping.py

(Closes #239)

@ajassani ajassani self-requested a review July 24, 2025 20:18
@ajassani ajassani merged commit 433d09e into main Jul 24, 2025
1 check passed
@ajassani ajassani deleted the spanmore/add_aiter_fa_perf_model branch July 24, 2025 20:19
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.

Add aiter flash attention to perf model

2 participants