ggml-webgpu: FlashAttention refactor + standardize quantization support

[reeselevine]

Overview

With three separate FlashAttention paths depending on sequence length and device capability, the code was getting messy. Quantized KV-caches also weren't supported by the tile path, which means that quantized KV-caches wouldn't run in WebGPU in the browser. This PR does a number of refactors to clean up the paths and add the same quantized KV-cache functionality everywhere:

Support and Setup Refactors

In ggml-webgpu.cpp:

• supports_op: checks only whether the sg_matrix or tile shader paths will work. This is because the auto FlashAttention setting uses a sequence length of 1 to probe support, but we want to ensure that FlashAttention will also work for larger sequence lengths, e.g., during prefill. Otherwise, we may end up in scenarios where the FlashAttention tensor used at runtime (with a larger sequence length then the initial check) can't fit on the GPU and runs on the CPU instead, which would be slower than not using FlashAttention to begin with.
• get_alloc_size: First checks if the tensor will run using the vec path, because only that path needs extra buffer space. No longer relies on the full decisions code logic for the FlashAttention shaders.
• during graph_compute: Does a common setup (ggml_webgpu_flash_attn_prepare), and then splits into vec/non-vec specific paths, instead of mixing these together.

In ggml-webgpu-shader-lib.cpp:

• Splits the vec/non-vec paths into separate cached pipelines and creation functions, since the combined path ended up mixing information and required a bunch of conditionals.
• Separates out functions required for supports_op and splitting into vec/non-vec paths, so that they can be called from ggml-webgpu.cpp without also dragging along a bunch of decision information that isn't needed. This also allows the decision information to be computed once when needed during pipeline creation.

In the shaders:

• flash_attn_quant_staging.tmpl is added, which includes common helpers and dequantization to shared memory that is used by all of the FlashAttention shaders.
• quant_inner_loops.tmpl is added, which abstracts some of the common dequantization code between mul_mat and flash_attn. This format can be followed for future sharing of code between the two operations too.

New functionality

• pre_wgsl.hpp is updated to allow multi-line macros, used for the shared dequantization to shared memory in the FlashAttention shaders.
• The K and V cache formats are no longer coupled. For example, f16 can be used for one and q8_0 for another, which wasn't possible previously.
• flash_attn_tile.wgsl, which is the long sequence length FlashAttention implementation used in the browser, now supports q4_0 and q8_0 KV-cache quantization.

I've verified that using quantized KV-caches works with this change on the command line and in Chrome on an Apple M3. Safari does not support subgroups, which are required for the WebGPU FlashAttention shaders.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: yes, for some of the mechanical refactors. Design was all done by me.

[reeselevine]

Any reviews from other WebGPU contributors are appreciated :)

@yomaytk @Constannnnnt @ArberSephirotheca

[yomaytk]

Checking only K type was enough before this PR, but with K/V
decoupled in this PR, I think that ggml_webgpu_flash_attn_kv_direct needs to check V type too.

For example, the following test case currently fails, but passes after adding V type checking:
test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {1, 1}, 256, 1, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0));

[reeselevine]

yep good catch, fixed this and added that test case too

[yomaytk]

This is not related to this PR change, but is std::max needed here?

[reeselevine]

nope, removed

[yomaytk]

Sharing the logic between multiple flash_attn wgsl files is nice! Also, the multi-line macro support would be useful for several cases. But any reason to use macros instead of functions here? Personally I'd prefer functions, but wdyt?

[reeselevine]

I was worried there would be a lot of arguments, but turns out that's not the case, so converted to functions. Was also able to remove some more code duplication.

I'll leave the multi-line macro ability in pre-wgsl though. Sometimes function calls, especially in hot inner loops, might reduce performance, so could be useful to use them in other places at some point.

[Constannnnnt]

Thanks! These changes look good to me.

[yomaytk]

good work, thanks! looks good to me.