facebookexperimental / facebookexperimental/triton

[AMD] Performant GFX950 support

Open
#1,084 1 comment 2 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.