ggml-cuda: MoE MMQ tile sizing + D=256 flash-attn tuning on RDNA3.5

[roberteg16]

ggml-cuda: MoE MMQ tile sizing + D=256 flash-attn tuning on RDNA3.5

Branch: rogarcia.mmq-moe-fattn (base: gfx11 @ 4b5ac96)
Target: gfx1151 (Radeon 8060S, Strix Halo, RDNA3.5)
Build: Release, GGML_HIP=ON, GGML_HIP_ROCWMMA_FATTN=OFF, GGML_HIP_ROOFLINE=OFF.

What this changes

Three prefill kernel-selection/config fixes for the HIP backend. All change which kernel/tile/launch is used, not the math, so outputs are bit-identical to the baseline.

1. Size the MMQ tile to MoE tokens-per-expert (mmq.cuh). For MUL_MAT_ID the tile width was minimized over ncols_max (worst case = all tokens routed to one expert), landing on mmq_x=128. In practice each expert receives ~ncols_dst/nchannels_y tokens (e.g. ~16 at pp512), so each 128-wide tile was ~1/8 full and the weight load dominated. The selection now derives the tile width from ~2x the average tokens-per-expert; batch-adaptive (large batches converge back to 128).

2. Route D=256 GQA attention to the mma-f16 kernel (fattn.cu, fattn-mma-f16.cuh). Raises the AMD_WMMA DKQ cap from 128 to 256 and routes D=256 to the mma-f16 kernel (packing ncols2 GQA heads into a full 16-wide tile). Decode (Q->ne[1] <= 2) still uses the vector kernel. NOTE: only takes effect in ROCWMMA_FATTN=ON builds - see the flash-attn section below.

3. Tune the D=256 tile flash-attn config for RDNA3.5 (fattn-tile.cuh). With ROCWMMA_FATTN=OFF (the fastest and intended build), D=256 prefill runs on the tile kernel. Adds an RDNA3.5-specific config that overrides the D=256, ncols=32 row (nbatch_K 128->64, occupancy 3->4); all other head sizes / ncols fall back to the shared RDNA table. Mirrors the MMQ device-table pattern; selection is matched on host (GGML_CUDA_CC_IS_RDNA3_5) and device (RDNA3_5 macro) so the __launch_bounds__ stay consistent with the host launch params.

Measured results (llama-bench, -fa 1, f16 KV, -ngl 999, -r 5, ROCWMMA_FATTN=OFF)

A/B by hot-swapping libggml-hip.so (gfx11 baseline vs this branch), interleaved per model (baseline run first / slightly cooler, so gains are conservative).

model	test	gfx11 (t/s)	this branch (t/s)	delta
Qwen3.6-35B-A3B	pp128	507.7 +/- 28.5	1054.1 +/- 12.8	+107.6%
Qwen3.6-35B-A3B	pp4096	1008.8 +/- 3.6	1424.0 +/- 34.2	+41.2%
Qwen3.6-35B-A3B	tg128	53.5 +/- 0.1	53.5 +/- 0.1	-0.0%
Qwen3.5-9B	pp4096	1074.0 +/- 17.9	1069.7 +/- 16.4	-0.4%
Qwen3.5-9B	tg128	35.7 +/- 0.0	35.7 +/- 0.0	+0.1%
Qwen3.5-4B	pp128	1729.2 +/- 166.7	1745.3 +/- 134.4	+0.9%
Qwen3.5-4B	pp4096	1808.8 +/- 68.2	1802.7 +/- 65.8	-0.3%
Qwen3.5-4B	tg128	58.7 +/- 0.3	58.6 +/- 0.3	-0.1%
Qwen2.5-7B	pp128	1381.6 +/- 76.4	1384.6 +/- 64.4	+0.2%
Qwen2.5-7B	pp4096	1419.1 +/- 46.3	1419.0 +/- 49.4	-0.0%
Qwen2.5-7B	tg128	45.5 +/- 0.1	45.5 +/- 0.1	+0.0%
llama-2-7b	pp128	1252.9 +/- 46.2	1252.5 +/- 46.0	-0.0%
llama-2-7b	pp4096	1161.5 +/- 26.1	1159.9 +/- 24.9	-0.1%
llama-2-7b	tg128	46.1 +/- 0.1	46.1 +/- 0.2	-0.1%
Qwen2.5-0.5B	pp128	8107.8 +/- 1305	8162.4 +/- 1183	+0.7%
Qwen2.5-0.5B	pp4096	12874.8 +/- 47.0	12892.0 +/- 17.6	+0.1%
Qwen2.5-0.5B	tg128	314.1 +/- 1.0	314.3 +/- 1.1	+0.1%

End-to-end latency (ms; derived tokens / t/s). e2e* = 4096 prefill + 128 decode

model	pp128 (base/br)	pp4096 (base/br)	tg128 (base/br)	e2e* (base/br)
Qwen3.6-35B-A3B	252 / 121	4060 / 2876	2391 / 2391	6451 / 5268 (-18%)
Qwen3.5-9B	135 / 120	3814 / 3829	3589 / 3586	7402 / 7415 (+0%)
Qwen3.5-4B	74 / 73	2264 / 2272	2181 / 2183	4445 / 4455 (+0%)
Qwen2.5-7B	93 / 92	2886 / 2887	2812 / 2810	5698 / 5697 (-0%)
llama-2-7b	102 / 102	3527 / 3531	2776 / 2778	6303 / 6310 (+0%)
Qwen2.5-0.5B	16 / 16	318 / 318	408 / 407	726 / 725 (-0%)

TTFT at context depth 4096 (Qwen3.6-35B-A3B, -p 4096 -n 128 -d 4096 -r 10)

lib	pp4096@d4096 (t/s)	TTFT
gfx11 (baseline)	916.6	4469 ms
this branch	1235.9	3314 ms (-25.8%)

Reading the results

• The MoE MMQ tile sizing is the real, durable win: the 35B-A3B MoE gains +108% at pp128, +41% at pp4096, and cuts full-request (4096+128) latency by 18%. Decode is unchanged (prefill-only effect).
• The D=256 flash-attn routing is inert with ROCWMMA_FATTN=OFF (it lives behind a gate that is compiled out), so it shows ~0% in the tables above. It is a large win in ROCWMMA_FATTN=ON builds, growing with context length - see the dedicated section below.
• Non-D256 dense models (Qwen2.5, llama-2-7b) are within noise, as expected.

Flash-attn commit: motivating config is GGML_HIP_ROCWMMA_FATTN=ON

The D=256 routing lives inside the ggml_cuda_should_use_wmma_fattn(cc) branch of ggml_cuda_get_best_fattn_kernel. That helper returns false whenever GGML_HIP_ROCWMMA_FATTN is not compiled in, so with ROCWMMA_FATTN=OFF the whole branch (and this commit's routing) is never reached - which is why the tables above show the FA change as a no-op. The commit exists to fix the ON build: there, RDNA3.5 D=256 GQA attention is dispatched to the rocWMMA kernel, which is badly underutilized, and this reroutes it to the mma-f16 kernel.

To prove the flash-attn commit helps on its own (independent of the MoE tile fix), an FA-only lib was built (gfx11 + only commit 47beebe, mmq.cuh untouched) and swept against depth on Qwen3.6-35B-A3B, ROCWMMA_FATTN=ON, -p 512 -n 0 -r 5. Attention's share of prefill grows with context depth, so the D=256 kernel choice dominates as depth increases:

config	gfx11 (t/s)	FA-only (t/s)	FA-only vs gfx11	full branch (t/s)	branch vs gfx11
p512 d0	1040.5	1133.0	+8.9%	1581.5	+52.0%
p512 d4096	602.7	992.0	+64.6%	1368.6	+127.1%
p512 d16384	257.1	699.5	+172.1%	863.2	+235.8%
p512 d32768	145.2	508.6	+250.3%	587.1	+304.4%
p4096 d4096	505.4	907.4	+79.5%	1224.6	+142.3%

TTFT at the p4096 d4096 point (prompt = 4096 / t/s x 1000):

lib	pp4096@d4096 (t/s)	TTFT
gfx11 (ON)	505.4	8104 ms
FA-only (ON)	907.4	4514 ms (-44% from FA alone)
full branch (ON)	1224.6	3345 ms (-59%)

So the flash-attn commit is a large, clean win within ROCWMMA_FATTN=ON builds and scales with context length. It is inert (harmless) in ROCWMMA_FATTN=OFF.

But ON + this commit is still not faster than OFF

Comparing the branch (FA commit present in both) across build flags on Qwen3.6-35B-A3B shows ON is equal-to-slower than OFF, progressively so at depth:

config	ON (t/s)	OFF (t/s)	ON vs OFF
p512 d0	1581.5	1593.8	-0.8%
p512 d4096	1368.6	1409.5	-2.9%
p512 d16384	863.2	942.9	-8.4%
p512 d32768	587.1	691.2	-15.1%
p4096 d4096	1224.6	1238.1	-1.1%

Reason (ggml_cuda_get_best_fattn_kernel, gfx1151): with OFF, D=256 skips the wmma branch and falls through to BEST_FATTN_KERNEL_TILE; with ON, this commit selects mma-f16. On this APU the tile kernel beats mma-f16 for D=256, more so at long context. Ordering: OFF (tile) > ON + this commit (mma-f16) >> ON baseline (rocWMMA).

RDNA3.5 tile flash-attn config (the OFF-path win)

Since ROCWMMA_FATTN=OFF is the fastest build, the highest-value lever is the tile kernel it uses for D=256 prefill. Its per-arch config table (fattn-tile.cuh, ggml_cuda_fattn_tile_get_config_amd_rdna) is shared across all RDNA GPUs and was not tuned for gfx1151. For a D=256 long prefill the launcher lands on the ncols=32 row, whose fields are (nthreads, occupancy, nbatch_fa, nbatch_K).

Swept on Qwen3.6-35B-A3B (-p 512 -n 0, deep context where attention dominates): nbatch_K 128->64 and occupancy 3->4 are the wins; nbatch_fa=64 and nthreads=256 are already optimal. Confirmed A/B (interleaved, upstream RDNA row 256,3,64,128 vs tuned 256,4,64,64):

depth (p512)	upstream RDNA (t/s)	tuned RDNA3.5 (t/s)	delta
d4096	1332.7	1358.3	+1.9%
d16384	911.7	967.9	+6.2%
d32768	673.6	721.4	+7.1%

The gain scales with context length (attention-bound) and is ~0 at shallow context - no regression. Verified again after wiring the RDNA3.5 branch (default build): +7.1% d16384, +6.0% d32768. Only D=256/ncols=32 on RDNA3.5 changes; every other GPU, head size, and ncols falls back unchanged. Output is bit-identical (launch/tiling only). Scope left for follow-up: the D=256 ncols=16/8 rows (short prompts / other GQA packings) still fall back to the shared RDNA values.

Recommendation

• The MoE MMQ tile commit is the clear keeper - it wins in both build configs.
• The RDNA3.5 tile-config commit improves the fastest (OFF) path itself, +6-7% on long-context D=256 prefill, with no downside - keep it.
• The flash-attn (mma-f16) commit only rescues ROCWMMA_FATTN=ON builds from the rocWMMA kernel; it does not make ON beat OFF. If the shipped build is OFF, it is a safety net at best - keep it solely to protect ON builds, but do not sell it as a net speedup over the OFF default.

Caveats

• gfx1151 is a thermally throttling APU; ppN t/s swing run-to-run (see large pp128 stddevs). Decode (tgN) is stable to ~1%.
• Correctness: tile width tiles the token dim (not K), so MoE output is bit-identical; PPL reported unchanged within error in the roadmap.
• Models benchmarked are the available Q4_K_M subset on this box (Qwen3.5-* stand in for the requested Qwen3-*); several requested models were absent.

[mgehre-amd]

Nice wins! Can we split the PR to have each of the changes in its own PR?