mha api refactor by slippedJim · Pull Request #1573 · ROCm/aiter

slippedJim · 2025-12-05T10:13:18Z

Motivation

Technical Details

Upload ASM pre/post processing kernels
Use unified params on all bwd kernels
Refine ASM kernel API

Test Plan

Test Result

Passed smoke test on gfx942 & gfx950
Passed python test on gfx942 & gfx950

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Copilot

Pull request overview

This PR refactors the Multi-Head Attention (MHA) backward pass API, consolidating arguments into a unified structure and simplifying the kernel selection and invocation process.

Key Changes

Unified MHA backward arguments into a single mha_bwd_args structure that encapsulates stream configuration, data types, and all tensor pointers
Introduced new CSV-based kernel configuration files for ASM kernel selection (odo, dqdkdv, dq_convert)
Added new C++ implementation for v3 backward pass with heuristic kernel selection

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
op_tests/cpp/mha/compile.py	Updated code generation commands for backward pass compilation
op_tests/cpp/mha/benchmark_mha_bwd.cpp	Simplified mha_bwd function call with new unified args structure
hsa/gfx942/fmha_v3_fwd/toy_format.hpp	Added custom string formatting utility for kernel name generation
hsa/gfx942/fmha_v3_fwd/fmha_v3_fwd.cpp	Added forward pass kernel selector with file/kernel name generation
hsa/gfx942/fmha_v3_fwd/aiter_hip_common.h	Added common HIP utilities and ASM kernel wrapper classes
hsa/gfx942/fmha_v3_bwd/*.csv	Added kernel configuration tables for backward pass kernels
csrc/include/mha_bwd.h	Refactored API with new args structure and detailed documentation
csrc/cpp_itfs/fmha_v3_bwd.cpp	Implemented new backward pass with heuristic kernel selection logic
aiter/jit/optCompilerConfig.json	Updated libmha_bwd configuration to include new source file

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Copilot

Pull request overview

Copilot reviewed 19 out of 63 changed files in this pull request and generated 7 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

* update bwd api * finish gfx942: need update dqdkdv kernel * update gfx950 csv * Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * remove expired files * update co file * update recompiled kernels && fix recompile kernel dispatch * fix compile config * update gfx942 pre-post kernel * fix asm varlen dkv reduction * remove log * fix swa on gfx942 * format file * remove not used files * add early quit condition for fmha_v3_bwd --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: jingchao <jichen@amd.com>

Summary: Pull Request resolved: pytorch#172592 update ck to commit ROCm/composable_kernel@292df27 update aiter to commit https://github.com/ROCm/aiter/commits/d3a4a0f1f5c83ee104a612e02447d03b0e38a45b/ changes of caffe2/aten/src/ATen/CMakeLists.txt and caffe2/caffe2/CMakeLists.txt are adopted from pytorch#161759 updated caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp to match the ck version in https://github.com/ROCm/composable_kernel/blob/292df2719f28cd01464d5d059820684790c101da/include/ck_tile/host/kernel_launch.hpp update aiter fav3 bwd codegen according to changes in ROCm/aiter#1573 update caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck mha fwd/bwd kernels according to the interfaces in https://github.com/ROCm/composable_kernel/tree/292df2719f28cd01464d5d059820684790c101da/example/ck_tile/01_fmha Differential Revision: D88991877

Summary: Pull Request resolved: #172592 update ck to commit ROCm/composable_kernel@292df27 update aiter to commit https://github.com/ROCm/aiter/commits/d3a4a0f1f5c83ee104a612e02447d03b0e38a45b/ changes of caffe2/aten/src/ATen/CMakeLists.txt and caffe2/caffe2/CMakeLists.txt are adopted from #161759 updated caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp to match the ck version in https://github.com/ROCm/composable_kernel/blob/292df2719f28cd01464d5d059820684790c101da/include/ck_tile/host/kernel_launch.hpp update aiter fav3 bwd codegen according to changes in ROCm/aiter#1573 update caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck mha fwd/bwd kernels according to the interfaces in https://github.com/ROCm/composable_kernel/tree/292df2719f28cd01464d5d059820684790c101da/example/ck_tile/01_fmha Reviewed By: iseeyuan Differential Revision: D88991877

Summary: Pull Request resolved: pytorch#172592 update ck to commit ROCm/composable_kernel@292df27 update aiter to commit https://github.com/ROCm/aiter/commits/d3a4a0f1f5c83ee104a612e02447d03b0e38a45b/ changes of caffe2/aten/src/ATen/CMakeLists.txt and caffe2/caffe2/CMakeLists.txt are adopted from pytorch#161759 updated caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp to match the ck version in https://github.com/ROCm/composable_kernel/blob/292df2719f28cd01464d5d059820684790c101da/include/ck_tile/host/kernel_launch.hpp update aiter fav3 bwd codegen according to changes in ROCm/aiter#1573 update caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck mha fwd/bwd kernels according to the interfaces in https://github.com/ROCm/composable_kernel/tree/292df2719f28cd01464d5d059820684790c101da/example/ck_tile/01_fmha Reviewed By: iseeyuan Differential Revision: D88991877

Summary: update ck to commit ROCm/composable_kernel@fcc9372 update aiter to commit ROCm/aiter@9a469a6 changes of caffe2/aten/src/ATen/CMakeLists.txt and caffe2/caffe2/CMakeLists.txt are adopted from pytorch#161759 updated caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp to match the ck version in https://github.com/ROCm/composable_kernel/blob/292df2719f28cd01464d5d059820684790c101da/include/ck_tile/host/kernel_launch.hpp update aiter fav3 bwd codegen according to changes in ROCm/aiter#1573 update caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck mha fwd/bwd kernels according to the interfaces in https://github.com/ROCm/composable_kernel/tree/292df2719f28cd01464d5d059820684790c101da/example/ck_tile/01_fmha Differential Revision: D88991877 Pull Request resolved: pytorch#172592 Approved by: https://github.com/alugorey, https://github.com/izaitsevfb

* update bwd api * finish gfx942: need update dqdkdv kernel * update gfx950 csv * Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * remove expired files * update co file * update recompiled kernels && fix recompile kernel dispatch * fix compile config * update gfx942 pre-post kernel * fix asm varlen dkv reduction * remove log * fix swa on gfx942 * format file * remove not used files * add early quit condition for fmha_v3_bwd --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: jingchao <jichen@amd.com>

* CK - gfx1250 support (#5) * Enable ROCM_CK_SDPA build * [submodule] composable_kernel and aiter update (pytorch#172592) Summary: update ck to commit ROCm/composable_kernel@fcc9372 update aiter to commit ROCm/aiter@9a469a6 changes of caffe2/aten/src/ATen/CMakeLists.txt and caffe2/caffe2/CMakeLists.txt are adopted from pytorch#161759 updated caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp to match the ck version in https://github.com/ROCm/composable_kernel/blob/292df2719f28cd01464d5d059820684790c101da/include/ck_tile/host/kernel_launch.hpp update aiter fav3 bwd codegen according to changes in ROCm/aiter#1573 update caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck mha fwd/bwd kernels according to the interfaces in https://github.com/ROCm/composable_kernel/tree/292df2719f28cd01464d5d059820684790c101da/example/ck_tile/01_fmha Differential Revision: D88991877 Pull Request resolved: pytorch#172592 Approved by: https://github.com/alugorey, https://github.com/izaitsevfb * Added MI450 supports and packages * Fix misalinged ck api * Replace aiter with ck for bwd * [ROCm] Bump AOTriton to 0.11.2b (pytorch#174105) Notable new features: * AOTriton 0.11.2b adds gfx1151/1152/1153 support. * Add precompiled AOTriton runtime for ROCM 7.2 * Match the sliding window attention behavior of `_flash_attention_forward/backward` with CUTLASS backend. Bug fixes: * Fixes pytorch#173204. Now all tests in `test/test_varlen_attention.py` are enabled on ROCm Notes: This replaces PR pytorch#173820 and pytorch#173469 Pull Request resolved: pytorch#174105 Approved by: https://github.com/jeffdaily * Fix philox data types for this version of ck * Update CK to use new gfx1250_pytorch branch * Add new gfx1250 compile flags for CK * add --targets to generate and a couple new compile flags * Remove default USE_ROCM_CK_SDPA --------- Co-authored-by: blorange-amd <bo.li2@amd.com> Co-authored-by: Yu Guo <yuguo@meta.com> Co-authored-by: Xinya Zhang <Xinya.Zhang@amd.com> * Updated aiter module * Fixed merged error * Fixed additional merged error * Reset USE_ROCM_CK_SDPA config --------- Co-authored-by: LugoReyes, Andy <Andy.LugoReyes@amd.com> Co-authored-by: Yu Guo <yuguo@meta.com> Co-authored-by: Xinya Zhang <Xinya.Zhang@amd.com>

slippedJim added 5 commits December 3, 2025 17:22

update

d7bfb21

update

d8766ba

update

4fc7983

update bwd api

a202ea3

update

5adbf89

Copilot AI review requested due to automatic review settings December 5, 2025 10:13

slippedJim marked this pull request as draft December 5, 2025 10:14

Copilot AI reviewed Dec 5, 2025

View reviewed changes

slippedJim and others added 11 commits December 8, 2025 16:52

finish gfx942: need update dqdkdv kernel

49748a8

update gfx950 csv

3149820

update

8974711

Merge branch 'main' into jim/dev/mha_refactor

2d658b1

update

00a892d

update

b9883fb

update

99561bd

update

74e2bb2

update

4d04248

update

4cd8f3a

Apply suggestion from @Copilot

5f0cb43

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

slippedJim requested a review from Copilot December 10, 2025 09:57

Copilot AI reviewed Dec 10, 2025

View reviewed changes

slippedJim added 5 commits December 11, 2025 05:54

update

79e2709

update

342cb2b

update

678a945

update

a8189c3

update

3dac41d

slippedJim marked this pull request as ready for review December 12, 2025 05:25

slippedJim requested a review from a team December 12, 2025 05:25

slippedJim and others added 2 commits December 12, 2025 13:27

remove expired files

e60fea6

update co file

56b6f7d

slippedJim added 3 commits December 15, 2025 22:38

fix swa on gfx942

b9fc1a0

format file

5e8222a

remove not used files

f9162fe

slippedJim requested review from danyao12 and valarLip December 16, 2025 00:58

slippedJim added 2 commits December 16, 2025 09:12

add early quit condition for fmha_v3_bwd

b222c5a

Merge branch 'main' into jim/dev/mha_refactor

9a59c3a

valarLip approved these changes Dec 16, 2025

View reviewed changes

valarLip self-assigned this Dec 16, 2025

slippedJim merged commit d0c313d into main Dec 16, 2025
22 checks passed

slippedJim deleted the jim/dev/mha_refactor branch December 16, 2025 14:10

zhuyuhua-v mentioned this pull request Dec 17, 2025

[dev/perf][Update] update dev/perf branch to main branch commit d0c313d #1660

Merged

yuguo68 mentioned this pull request Jan 15, 2026

[submodule] composable_kernel and aiter update pytorch/pytorch#172592

Closed

Navigation Menu

mha api refactor - #1573

mha api refactor#1573
slippedJim merged 34 commits into
mainfrom
jim/dev/mha_refactor

slippedJim commented Dec 5, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

slippedJim commented Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Key Changes

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

slippedJim commented Dec 5, 2025 •

edited

Loading