facebookexperimental / facebookexperimental/triton
[AMD] Performant GFX950 support
- Dominant language
- MLIR
- Stars
- 195
- Forks
- 72
- Avg merge
- 14m
- Merged PRs (30d)
- 3
Description
Setting up a live document/tracker for tracking work that needs to be done on GFX950/MI350 on TLX to get optimal performance. We will update and add more work items as we develop and progress.
Currently the two kernels we are going to support is F16 GEMM and F16 MHA
Optimal Common:
- [ ] Implement warp pipelining support in TLX for PingPong
- [ ] Ensure paddedSharedLayout to minimize bank conflict when possible.
- [ ] Enable extract_slice on TLX for dot slicing to reduce register pressure and more scheduling flexibility.
- [ ] Layout Inference across loops to enable fully pipelined code(e.g local prefetch gemm, `a` and `b` layout cannot be determined in this [local prefetch kernel](https://github.com/raikonenfnu/tlx_triton/blob/raikonenfnu/exp_gfx950/third_party/tlx/tutorials/amd_gemm_v5.py#L90-L105).
- [ ] Port LLIRSched from https://github.com/ROCm/triton/tree/matmul_4waves
Optimal GEMM:
- [ ] Implement GFX950 global prefetch kernel
- [ ] Implement GFX950 fully pipelined/local prefetch kernel (dependent on layout inference and similar to [this kernel ](https://github.com/raikonenfnu/tlx_triton/blob/raikonenfnu/exp_gfx950/third_party/tlx/tutorials/amd_gemm_v5.py#L90-L105)but need to get it to work and performant)
- [ ] Add dot slicing into pipelined gemm kernel
- [ ] Add loop unrolling to eliminate pointer rotations.
- [ ] Add L2 cache XCD remap for GFX950 gemm kernels
Optimal FA:
- [ ] Port LLIRSched from https://github.com/ROCm/triton/tree/matmul_4waves
- [ ] Implement naive FA on gfx950
- [ ] Implement pipelined FA on gfx950 ([similar gluon kernel](https://github.com/raikonenfnu/ml-kernels/blob/main/cdna4/fa/f16_fa_gfx950.py))
- [ ] Implement FAv3 on gfx950 (dependent on pipelined FA + pipelining)
- [ ] Add L2 cache XCD remap for GFX950 FA kernels
Contributor guide
Assessment
This issue has not been assessed yet.