[FA] Improve performance of shapes <95% on advanced path - 32x32x512, 4x32x4096, 2x32x8192

[whitneywhtsang]

There are currently 3 shapes (causal=false, d_head=64) that have performance <95% of XeTLA.

• 32, 32, 512: 78%
• 4, 32, 4096: 82%
• 2, 32, 8192: 93%

[quintinwang5]

Huge SBIDStalls in Triton kernel.


Also more DistStalls in Triton kernel(the red line)


WIP to investigate and finetune.

[quintinwang5]

Status: Did lots of finetunes and experiments. Most finetune configs are harmful to the performance. Modified llir by hand and execute it, there is also no big performance gains(ex. While dpas insts are atomic, we'll have more spills. It will offset the benefit from atomic dpas). The reason is higher SBIDStall and DistStall for load/sync insts than XETLA. This shows we are in short of registers. We may need helps from XETLA team or others.

[quintinwang5]

Cache hit of XETLA:

Cache hit of Triton:

Triton with PR:

[quintinwang5]

Overall status:

[quintinwang5]

According to the new overall status:

• 32, 32, 512 , 64, False: 78% --> 92.99%
• 4, 32, 4096, 64, False: 82% --> 96.17% (not by this issue)
• 2, 32, 8192, 64, Flase: 93% --> 97.85 (not by this issue)

extra N_CTX=512 shapes:

• 32, 16, 512, 128, False: 66.41% --> 80.91%
• 32, 16, 512, 128, True: 60.10% --> 80.26%
• 32, 32, 512, 64, True: 65.91% --> 88.44%

Still need some works to bring 92.99% up to 95%+.