Skip to content

ggml-cuda: use mma-f16 flash-attn for D=256 on RDNA3.5#40

Open
roberteg16 wants to merge 1 commit into
gfx11from
rogarcia.fattn-mma-f16-d256
Open

ggml-cuda: use mma-f16 flash-attn for D=256 on RDNA3.5#40
roberteg16 wants to merge 1 commit into
gfx11from
rogarcia.fattn-mma-f16-d256

Conversation

@roberteg16

Copy link
Copy Markdown

What this changes

A prefill optimization for D=256 GQA attention on RDNA3.5 when built with rocWMMA FlashAttention (-DGGML_HIP_ROCWMMA_FATTN=ON). On this path ggml_cuda_get_best_fattn_kernel would pick the wmma-f16 kernel, which for GQA leaves its 16-wide tile mostly empty. The mma-f16 kernel also lowers to RDNA WMMA instructions here, but it packs ncols2 GQA query heads into a full 16-wide tile, so it uses the same hardware far more efficiently.

Two changes: (1) in fattn.cu, route D=256 to BEST_FATTN_KERNEL_MMA_F16 when the GQA optimization applies and the tile is actually filled (gqa_opt_applies && Q->ne[1]*gqa_ratio_eff >= 16); (2) in fattn-mma-f16.cuh, raise the AMD_WMMA path's head-dim guard from DKQ > 128 to DKQ > 256 so the D=256 kernel is allowed to compile. Both are restricted to RDNA3.5 (GGML_CUDA_CC_IS_RDNA3_5 on the host, the RDNA3_5 macro for the device-side limit), the only architecture this was tuned and validated on; other AMD WMMA architectures keep the previous DKQ <= 128 behavior.

Benchmarks

Measured on gfx1151 (Radeon 8060S), Qwen3.6-35B-A3B Q4_K_M, -ngl 999 -r 1, built with -DGGML_HIP_ROCWMMA_FATTN=ON. Baseline is gfx11 at the same commit base, also built with rocWMMA on.

test baseline t/s this PR t/s delta
pp128 563.90 556.10 -1.4% (noise)
pp1024 993.58 1129.89 +13.7%
tg128 52.99 52.95 -0.1% (noise)

Larger-batch prefill (pp1024) improves clearly; pp128 and decode are within run-to-run noise. This only affects the rocWMMA-on build - with rocWMMA off, ggml_cuda_should_use_wmma_fattn returns false and the routing change is never reached.

Route D=256 GQA attention to the mma-f16 kernel and allow its AMD_WMMA path
up to DKQ=256. On RDNA the kernel still runs on WMMA instructions, and it
packs ncols2 GQA heads into a full 16-wide tile that the wmma kernel leaves
mostly empty.

Assisted-by: Claude Opus 4 (1M context)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant