Adding gfx908 support

[UD-mmcminn]

Motivation

I use MI100's for AI development and it would be nice if we could keep support for those. There does not appear to be anything in aotriton that blocks MI100 support.

Technical Details

Just added gfx908 to the places that need it.

Test Plan

I tried testing, but was unable to get the pytest suite to work on my test system. If you can provide some info on that I'd be happy to test.

Test Result

N/A

Submission Checklist

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

[UD-mmcminn]

I have a box with 8 MI100s with infinity links if testing is needed. I just require some support getting the tests to run. I spent an hour mucking around but it just would not cooperate. I'm likely missing something silly, but I just can't seem to identify what. Trying to replicate the steps in the scripts under .ci has not worked. I suspect the python venv is not playing nice with the .so but I don't see a way to force the install into the venv.

[xinyazhang]

The first step to enable gfx908 is to test SDPA/FA kernel works with Triton.
There is no magic in aotriton that can drastically support architectures that Triton can't.

[UD-mmcminn]

Gotcha, the issue is, this limitation is breaking pytorch for gfx908 right now. I've got a fork and a patched version of torch I am working with to try and get a test working. Having a weird error with hip and the MATH backend I am tracking down. Likely a flags error. Anyway, maybe experimentally enable it while we get it figured out and debugged?

[xinyazhang]

Anyway, maybe experimentally enable it while we get it figured out and debugged?

The problem is the capability to compile the kernel for gfx908 is questionable. IIRC there is no gfx908 MFMA support in Triton compiler, and its MFMA hardware only works with fp16 datatype.

[UD-mmcminn]

GFX908(MI100) does support BF16 with the MFMA hardware. I know this 100% since I've been patching things to use it. AMD has a bad habit of dropping support / ignoring prior gen cards. Anyway, the BF16 MFMA hardware is half width but it just requires 2 instructions instead of 1. That said, CK and various underlying libraries already do this. I'll do some additional reasearch into triton but as far as I can tell, the specific MFMA code in the underlying stuff is the MI300 series stuff. I'm still wrapping my head around the internals of the project.

Also aotrition does compile and appears to work for GFX908 after applying this PR. I just need to find a way to validate that. The other PR has that script but I'm getting an error with the MATH backend. The other two backends work. I think I didn't compile it correctly or something to work for CPU which seems very odd.

[xinyazhang]

• Tuning database of gfx908 is missing. All kernel will be compiled with default options. (It's okay if it's intentional)
• Changes to v2src/v2python can be removed. They are deprecated.
• Need to add gfx90a to isArchExperimentallySupported, unless someone volunteers to run full tuning/testing process on gfx908 before each release.

[UD-mmcminn]

• Tuning database of gfx908 is missing. All kernel will be compiled with default options. (It's okay if it's intentional)
• Changes to v2src/v2python can be removed. They are deprecated.
• Need to add gfx90a to isArchExperimentallySupported, unless someone volunteers to run full tuning/testing process on gfx908 before each release.

Roger that. I am working on these. Specifically the FLASH_ATTENTION implementation on the gfx908. The MATH backend appears to work but the FLASH_ATTENTION one isn't behaving and I am trying to nail down the cause.

[UD-mmcminn]

Tuning running... Builds timed out for the bwd_kernels with HEAD_DIM>=256 and as low as M >= 32 N >= 32. I'm assuming this is VGPR pressure. As for the tuning, it's looking like about a day to run? Does this seem correct?

[UD-mmcminn]

I have generated the tuning database, well most of it. There were some configs that would cause crashes when dropout was involved. Also, larger sequence would run out of vram. The JSON is somewhat large and I am unsure what to do with it.

[IMbackK]

@UD-mmcminn just curious if you got any further on this.