[TRITON] Add Positional Encoding (PE) support to Triton MHA kernels#1203

Merged

vgokhale merged 65 commits into355_wipfrom

cagri/bruno_changes

Oct 16, 2025

Commits on Oct 14, 2025

Added PE support for MHA

azaidy
authored and
cagrikymk
committed
Enable MHA backward tests

brunomazzottiamd
authored and
cagrikymk
committed
Apply `black` formatting

brunomazzottiamd
authored and
cagrikymk
committed
Add `PE_HEAD_DIM` and `ACTUAL_PE_HEAD_DIM` parameters to backward kernels

brunomazzottiamd
authored and
cagrikymk
committed
[WIP] Add `q_pe` and `k_pe` to `_bwd_dkdv_inner`

brunomazzottiamd
authored and
cagrikymk
committed
Draft PE backward test

brunomazzottiamd
authored and
cagrikymk
committed
Fix minor errors

brunomazzottiamd
authored and
cagrikymk
committed
Document test failures

brunomazzottiamd
authored and
cagrikymk
committed
Split `dk` into non-positional and positional parts

brunomazzottiamd
authored and
cagrikymk
committed
Disable PE tests by default

brunomazzottiamd
authored and
cagrikymk
committed
Fix compilation error in `_bwd_dkdv_inner`

brunomazzottiamd
authored and
cagrikymk
committed
Reduce diff content to facilitate later code review

brunomazzottiamd
authored and
cagrikymk
committed
Implement `dq` part of causal backward

brunomazzottiamd
authored and
cagrikymk
committed
Reduce test scope and document failures

brunomazzottiamd
authored and
cagrikymk
committed
Create 1st version of forward BSHD + PE test

brunomazzottiamd
authored and
cagrikymk
committed
Document BSHD backward failures

brunomazzottiamd
authored and
cagrikymk
committed
Document forward accuracy debugging attempts

brunomazzottiamd
authored and
cagrikymk
committed
Try smaller shapes to debug MHA forward with PE

brunomazzottiamd
authored and
cagrikymk
committed
Add code to dump tensors in fwd test

brunomazzottiamd
authored and
cagrikymk
committed
Disable fused bwd benchmarking

brunomazzottiamd
authored and
cagrikymk
committed
Dumping more tensors on fwd unit test

brunomazzottiamd
authored and
cagrikymk
committed
Document some debugging attempts

brunomazzottiamd
authored and
cagrikymk
committed
1st version of bwd fine tune script

brunomazzottiamd
authored and
cagrikymk
committed
hack pe fwd + bwd test to enable bwd shape tuning

brunomazzottiamd
authored and
cagrikymk
committed
minor improvements to bwd fine tune script

brunomazzottiamd
authored and
cagrikymk
committed
add new best config for pe bwd on MI350X

brunomazzottiamd
authored and
cagrikymk
committed
reduce tuning space

brunomazzottiamd
authored and
cagrikymk
committed
Revert "add new best config for pe bwd on MI350X"

brunomazzottiamd
authored and
cagrikymk
committed
add new best config for pe bwd on MI350X - correct compiler

brunomazzottiamd
authored and
cagrikymk
committed
Revert "add new best config for pe bwd on MI350X - correct compiler"

brunomazzottiamd
authored and
cagrikymk
committed
add new best config for pe bwd on MI350X - again

brunomazzottiamd
authored and
cagrikymk
committed
fix forward bug: k pe pointer was not advancing

brunomazzottiamd
authored and
cagrikymk
committed
find small shape for dq fail + dq tensor dump

brunomazzottiamd
authored and
cagrikymk
committed
cleanup attention ref implementation

brunomazzottiamd
authored and
cagrikymk
committed
[WIP] debugging wrong dq computation

brunomazzottiamd
authored and
cagrikymk
committed
fix dq bug in bwd

brunomazzottiamd
authored and
cagrikymk
committed
tweak bwd tuning script

brunomazzottiamd
authored and
cagrikymk
committed
group PE tests together

brunomazzottiamd
authored and
cagrikymk
committed
improve coverage of test_mha_with_pe

brunomazzottiamd
authored and
cagrikymk
committed
revert debugging hack

brunomazzottiamd
authored and
cagrikymk
committed
ommit linter check

brunomazzottiamd
authored and
cagrikymk
committed
improve coverage of test_mha_varlen_with_pe

brunomazzottiamd
authored and
cagrikymk
committed
fwd PE support only for power of 2 heads

brunomazzottiamd
authored and
cagrikymk
committed
[WIP] improve coverage of test_mha_backward_with_pe

brunomazzottiamd
authored and
cagrikymk
committed
revert debugging comments related to wrong dq calculation

brunomazzottiamd
authored and
cagrikymk
committed
[WIP] a lot of things...

brunomazzottiamd
authored and
cagrikymk
committed
add test_mha_backward_varlen_with_pe

brunomazzottiamd
authored and
cagrikymk
committed
enable more bwd tests

brunomazzottiamd
authored and
cagrikymk
committed
force "onekernel" bwd implementation in bwd varlen pe test

brunomazzottiamd
authored and
cagrikymk
committed
comment annoying print statement

brunomazzottiamd
authored and
cagrikymk
committed
update bwd tuning script

brunomazzottiamd
authored and
cagrikymk
committed
comment cleanup

brunomazzottiamd
authored and
cagrikymk
committed
distinguish PE and no-PE best configs

brunomazzottiamd
authored and
cagrikymk
committed
raise error on fused backward + PE inputs

brunomazzottiamd
authored and
cagrikymk
committed
replace print statement by aiter logger

brunomazzottiamd
authored and
cagrikymk
committed
add best bwd config for MI300

brunomazzottiamd
authored and
cagrikymk
committed
[TO BE TESTED] update benchmark script

brunomazzottiamd
authored and
cagrikymk
committed
[TO BE TESTED] create dirty script for fwd tuning

brunomazzottiamd
authored and
cagrikymk
committed
improve bench script

brunomazzottiamd
authored and
cagrikymk
committed
fix wrong assertion

brunomazzottiamd
authored and
cagrikymk
committed
fix dirty tuning scripts

brunomazzottiamd
authored and
cagrikymk
committed
add TODO comment to bench script

brunomazzottiamd
authored and
cagrikymk
committed
add new best config for MI300 fwd

brunomazzottiamd
authored and
cagrikymk
committed
remove constexpr seqlens
cagrikymk
committed

Commits on Oct 15, 2025

remove tuning files
cagrikymk
committed