Skip to content

ggml-webgpu: simplify flash_attn shaders and refactor several WGSL shaders#26134

Draft
yomaytk wants to merge 3 commits into
ggml-org:masterfrom
yomaytk:refactor-wgsl
Draft

ggml-webgpu: simplify flash_attn shaders and refactor several WGSL shaders#26134
yomaytk wants to merge 3 commits into
ggml-org:masterfrom
yomaytk:refactor-wgsl

Conversation

@yomaytk

@yomaytk yomaytk commented Jul 26, 2026

Copy link
Copy Markdown
Member

Overview

This PR includes the changes from #25956, so I'll open this after the PR is resolved.

This PR refactors several WGSL files and simplifies the flash attention WGSL files. The changes to each WGSL shader are as follows:

  • concat.wgsl, rms_norm_mul.wgsl, row_norm.wgsl, solve_tri.wgsl, ssm_scan.wgsl: Remove unused params.
  • conv2d.wgsl, conv2d_dw.wgsl, im2col.wgsl: Define the macros as the actual type for weights, input, and output buffer and remove some unnecessary functions.
  • soft_max.wgsl: Remove the unused param ne. Also, binding numbers for each tensor buffer are computed sequentially instead of defining buffers with nested macros, since the nested macro structures aren't easy to read.
  • flash_attn_decls.tmpl: This file contains the common definitions of params and types shared among flash attention WGSL shaders. Binding numbers for the tensor buffers are computed in the same way as in soft_max.wgsl.
  • flash_attn_staging.tmpl: flash_attn_quant_staging.tmpl was renamed to this file, and load_(k|v)_tile_block for the float type is extracted from the three flash attention shaders.
  • flash_attn.wgsl, flash_attn_tile.wgsl, flash_attn_vec_split.wgsl: Common code is extracted into flash_attn_decls.tmpl and flash_attn_staging.tmpl.

Requirements

  • I have read and agree with the contributing guidelines
  • AI usage disclosure: YES, I used AI to find the redundant code for the all wgsl shaders, and to refactor the several parts. Those changes are reviewed and the other parts (mainly for FA) were written by me.

@github-actions github-actions Bot added ggml changes relating to the ggml tensor library for machine learning WebGPU labels Jul 26, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ggml changes relating to the ggml tensor library for machine learning WebGPU

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant