Placed MMA gateway over WMMA gateway

[adam-kabbara]

Overview

When the build parameter GGML_HIP_ROCWMMA_FATTN=ON is added, it forces all fast attention kernel pathways to BEST_FATTN_KERNEL_WMMA_F16 given a large enough batch size, and applicable head sizes (else we fall back to VEC or TILE). However, the BEST_FATTN_KERNEL_MMA_F16 path is faster for models with head sizes <=128 as seen in the two benchmark tests performed on RDNA3 and RDNA4 below.

This patch fixes two things in one.

1. By checking if the MMA path is applicable and trying it out before WMMA, performance for models with head sizes <= 128 improves.
2. Allows GGML_HIP_ROCWMMA_FATTN to follow its definition in CMakeLists.txt more tightly.

Additional information

Benchmarking using two llama-bench run lines

./llama-bench -m <model> -ngl 99 -pg 512,128 -fa 1
./llama-bench -m <model> -p 4096,8192 -n 128 -b 2048,4096 -ngl 99 -fa 1

Using the following build line for both before and after patch.

  -DGGML_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx1201 \
  -DCMAKE_PREFIX_PATH="$DEVEL" -DCMAKE_HIP_COMPILER="$DEVEL/lib/llvm/bin/amdclang++" \
  -DLLAMA_CURL=OFF -DCMAKE_BUILD_TYPE=Release -DGGML_VULKAN=OFF \
  -DGGML_HIP_ROCWMMA_FATTN=ON

Tables below show % difference in tok/s from pre-patch to post-patch.

RDNA4 Benchmark

GGML_HIP_ROCWMMA_FATTN Definition

In CMakeLists.txt it mentions

option(GGML_HIP_ROCWMMA_FATTN               "ggml: enable rocWMMA for FlashAttention"         OFF)

This build parameter should just enable rocWMMA, not force rocWMMA. Before this patch, if the above build parameter is turned on, it forces the fattn kernel path to WMMA given a large enough batch size. However, with this patch, the compiler choses either WMMA or MMA path depending on which is fast for the given model. If the build parameter is off, this patch will not have an effect.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure:

[ggml-gh-bot]

Hi @adam-kabbara, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[pwilkin]

Can you add -d 0,40000 to the bench? From my tests, GGML_HIP_ROCWMMA_FATTN=ON caused extreme prefill speed degradation at longer context sizes.