UPSTREAM PR #18610: ggml webgpu: initial flashattention implementation

[loci-dev]

Mirrored from ggml-org/llama.cpp#18610

This PR adds an initial version of FlashAttention2 in WebGPU. Along with the GPU code itself, this PR also adds a new preprocessor for WGSL shaders that should make it easier/less brittle to define new shaders going forwards. Details below:

Shader setup

• Most of the shaders right now are generated at build time using a relatively hacky Python script and template syntax that I wrote when doing initial development for the WebGPU backend. However, this probably won't continue to scale very well, especially with the number of options for FlashAttention, so I decided it was time to use a more general solution. However, there wasn't an existing preprocessor for WGSL that would work with C++ code, so I wrote one here: https://github.com/reeselevine/pre-wgsl. The preprocessor itself is one file, pre_wgsl.hpp, and should continue to track any changes/features added to the main preprocessor repository.
• To accommodate the various options for FlashAttention when compiling WGSL shaders, I added another new file to the WebGPU backend, ggml-webgpu-shader-lib.hpp, which generates the shader using a combination of structural parameters, e.g., head sizes, and performance parameters, like KV tile sizes. My idea is to expand this library approach use more sophisticated strategies going forwards and move other shaders over to use the preprocessor. From a performance perspective, shaders compile pretty fast, are mostly fixed for a given model, and are cached, so JIT compilation, at least for FlashAttention, seems to be the right call in my opinion.
• Some other minor changes to ggml-webgpu.cpp to handle the new FlashAttention code and JIT compilation.

FlashAttention shader itself

• For the most part this follows the FlashAttention2 paper, with a change for online softmax to make it subgroup-size agnostic. It also uses global KV loads if sizes are nicely divisible, since the KV tiles are not reused and pre-loading into shared memory really slows things down (at least on my M3).
• Even then, performance is not great right now (< 50% of the same Metal code). My testing shows that a lot of this slowdown basically boils down to the initial Q * K^T accumulation loop. I still need to do some debugging here, to figure out if the issue is something I'm doing wrong structurally, or if it's in the compilation from WGSL to Metal. This code also needs to be more thoroughly tested on other platforms. Perhaps someone who has written FlashAttention for one of the other backends could take a look at the shader and see if it looks reasonable, e.g., @jeffbolznv?
• Otherwise, this passes all the backend tests on my machine, so I think it's in a good state to merge as an initial implementation that can be improved upon as time goes on.

[loci-review]

Explore the complete analysis inside the Version Insights

Perfect! I've generated the summary report for your project. Here's what the analysis shows:

Key Findings:

✅ No Significant Performance Impact Detected

For the llama.cpp repository (Pull Request #821 by auroralabs-loci):

• Response Time: No modified functions showed changes greater than 2%
• Throughput: No modified functions showed changes greater than 2%

This indicates that the pull request is performance-neutral and safe to merge from a performance perspective. The changes appear to focus on functionality, bug fixes, or code quality improvements without affecting the performance characteristics of the codebase.

The comparison was made between:

• Base version: a7cb4ab1-e9f6-11f0-81f2-dbb430499cb5
• Target version: a6976c51-e9fb-11f0-81f2-dbb430499cb5