ggml-webgpu: tune subgroup split (d_split) in flash_attn_vec

[yomaytk]

Overview

This PR adds tuning of D_SPLIT (= subgroup_size / VEC_NE, where VEC_NE was the value hardcoded on master) in flash_attn_vec to split the subgroup when accumulating Q * K and P * V, based on head_dim and subgroup_size, instead of hardcoding it. The following table shows the TG performance improvement when the context length reaches 16K.

The command is like this: llama-bench -m Llama-3.2-3B-Instruct-Q4_K_M.gguf -fa 1 -p 0 -n 128 -d 16384 -r 3 -dev WebGPU

NVIDIA V100

Model	master (t/s)	PR (t/s)	Speedup
gemma4 E4B Q4_K_M	26.02 ± 0.16	27.21 ± 0.15	+4.6%
gpt-oss 20B MXFP4 MoE	26.35 ± 0.37	27.78 ± 0.36	+5.4%
llama 3B Q4_K_M	25.55 ± 1.44	34.78 ± 2.57	+36.1%

M2

Model	master (t/s)	PR (t/s)	Speedup
gemma4 E4B Q4_K_M	6.79 ± 1.53	6.83 ± 1.50	+0.6%
gpt-oss 20B MXFP4 MoE	20.12 ± 0.17	22.22 ± 0.16	+10.4%
llama 3B Q4_K_M	18.18 ± 0.05	20.18 ± 0.01	+11.0%

This tuning is similar to the Vulkan backend. In practice, VEC_NE is set to 8 in many cases, and I can confirm that it performs better than other values of master branch.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Yes, I used AI to investigate vulkan's flash attention implementation, and analyze the performance logs.

[yomaytk]

Vulkan sets this upper limit of D_SPLIT to 8u, but I can confirm that 4u works better on both M2 and V100.

[reeselevine]

pretty nice speedups!

[yomaytk]

@ggml-org/maintainers can I get a second review please?